bug in onnx export of aten_repeat_interleave_self_int results in incorrect code

[bas-aarts]

running this script:

import onnx
import torch

class MyModule(torch.nn.Module):
    def forward(self):
        return torch.repeat_interleave(torch.tensor([2]), 2, 0)

model = MyModule()
model.eval()
with torch.no_grad():
    torch.onnx.export(model, (), optimize=False).save("output.onnx")
onnx.shape_inference.infer_shapes_path("output.onnx", strict_mode=True)

results in the following error:

onnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] Inference error(s): (op_type:Identity, node name: node_repeat_interleave): [ShapeInferenceError] Inferred shape and existing shape differ in rank: (2) vs (1)

the issue is this line of code

onnxscript/onnxscript/function_libs/torch_lib/ops/core.py

Line 8314 in c72e810

if self_rank == 1:

.
This line checks the original rank, and not the current rank. It should therefore be:
if self_rank == 0:

versions of relevant packages:

onnx                          1.20.0
onnx-ir                       0.2.1
onnxscript                    0.7.0
torch                         2.11.0

[bas-aarts]

or, replace

onnxscript/onnxscript/function_libs/torch_lib/ops/core.py

Lines 8295 to 8298 in c72e810

	if dim is None:
	raise NotImplementedError("No conversion available yet when dim is None.")
	self_rank = len(self.shape)

with

    if dim is None:
        self = op.Reshape(self, [-1])
        dim = 0
        self_rank = 1
    else:
        self_rank = len(self.shape)

and remove

onnxscript/onnxscript/function_libs/torch_lib/ops/core.py

Lines 8314 to 8315 in c72e810

	if self_rank == 1:
	return op.Identity(tiled)

alltogether.

this fixes the issue, and adds support for the 'dim is None' case