Fix aten_roll for negative dims and shifts past the dimension length

[om singhal (Om-singhaI)]

torch.roll is circular. The shift is taken modulo the length of the dimension,
and a negative dims counts from the end. The lowering in
onnxscript/function_libs/torch_lib/ops/core.py does neither. It slices at
dim_size - shift for a positive shift and at -shift for a negative one, and
leans on Slice clamping to keep that in range, which only holds while the
shift stays inside a single wrap.

A negative dim emits a model onnxruntime will not load

_aten_roll_shift_and_dim_onnx reads the dimension length with
op.Shape(self, start=dim, end=dim + 1). For dim = -1 that is
start=-1, end=0, and ONNX reads end=0 as the absolute index 0, not as the
end of the shape. The range is empty, so Shape returns an empty tensor, and
that empty tensor is what reaches the ends input of Slice.

Unoptimized graph for torch.roll(x, shifts=1, dims=-1) on main, x of shape
(2, 3):

Shape    node_Shape_2   in=['x']                        out=['val_2']  attrs={'end': 0, 'start': -1}
Sub      node_Sub_4     in=['val_2', 'val_3']           out=['val_4']
Slice    node_Slice_6   in=['x', 'val_5', 'val_4', ...] out=['val_6']
Concat   node_roll      in=['val_9', 'val_6']           out=['roll']   attrs={'axis': -1}

torch.roll(x, shifts=1, dims=-1)
  torch        [[2.0, 0.0, 1.0], [5.0, 3.0, 4.0]]
  onnxruntime  [ONNXRuntimeError] : 1 : FAIL : Node (node_Slice_6) Op (Slice)
               [ShapeInferenceError] Incorrect or missing input value for starts and ends

Those nodes are this lowering rather than a torch decomposition. They are the
Shape, Sub, Slice, Slice, Concat chain from
_aten_roll_shift_and_dim_onnx, the output node keeps the node_roll name, and
editing core.py changes what onnxruntime returns, which is how every number
below was produced.

The export above uses optimize=False so the model gets as far as being saved.
With the default optimize=True it does not get that far: the empty tensor
trips the rewrite pass first and the export raises PassError. Same root cause.

Only dim = -1 breaks this way. dim = -2 on a rank 2 tensor gives
start=-2, end=-1, which is a nonempty range, and comes out correct.

A shift past the dimension length is silently wrong

x = torch.arange(6, dtype=torch.float32).reshape(2, 3)

torch.roll(x, shifts=-4, dims=1)
  torch        [[1.0, 2.0, 0.0], [4.0, 5.0, 3.0]]
  onnxruntime  [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]

torch.roll(x, shifts=7, dims=1)
  torch        [[2.0, 0.0, 1.0], [5.0, 3.0, 4.0]]
  onnxruntime  [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]

torch.roll(x, shifts=14)
  torch        [[4.0, 5.0, 0.0], [1.0, 2.0, 3.0]]
  onnxruntime  [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]

The boundary is narrow, which is why this went unnoticed. A shift of exactly one
dimension length is right because the rotation is the identity. A positive shift
between one and two lengths is also right, because dim_size - shift goes
negative and Slice reads a negative start as an index from the end, which
happens to be one wrap. It only goes wrong once a positive shift exceeds twice
the length, and for any negative shift whose magnitude is not a multiple of the
length. The last case above is the dims=() path, where the same thing happens
against the element count.

The fix

aten_roll normalizes a negative dim against the rank, which is known at export
time, in the same shape aten_unflatten already uses:

# PyTorch accepts negative dim as reversed counting
if dim < 0:
    dim = self_rank + dim

aten_roll_complex needs its own version. The real representation carries a
trailing axis for the real and imaginary parts, so a negative dim resolves
against a rank one larger than the one torch sees. This matches
aten_slice_complex and aten_squeeze_dim_complex:

if dim < 0:
    # Account for the complex dimension in ONNX
    dim = self_rank + dim - 1

Both helpers then take the shift modulo the length. ONNX Mod with fmod=0
gives the remainder the sign of the divisor, the same as Python, so
(-shift) % length is the split point for either sign of shift and any
magnitude:

dim_length = op.Shape(self, start=dim, end=dim + 1)
slice_length = op.Mod(op.Constant(value_ints=[-shift]), dim_length)

The branch on the sign of the shift goes away in both helpers. The second
Slice was passing the total element count as its ends and relying on
clamping; it now takes the dimension length it already has. The graph for the
case above goes from 11 nodes to 8, and Shape now reads start=1, end=2.

Testing

Added test_roll_wraps_shifts_and_normalizes_negative_dims and
test_roll_complex_wraps_shifts_and_normalizes_negative_dims to
tests/function_libs/torch_lib/e2e_ops_tests.py, following the optimize=False
pattern the tests around them use. Seven cases between them: negative dim,
positive and negative shifts past the dimension length, and the dims=() path
where the shift runs past the element count.

With only the core.py change reverted, all seven fail. Two fail at session
creation with the error quoted above, the other five with
AssertionError: Tensor-likes are not close!. With the fix all seven pass.

pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k roll
7 passed, 100 deselected in 6.88s

pytest tests/function_libs/torch_lib/ops_test.py -k roll
6 passed, 2 skipped, 1844 deselected, 102 subtests passed in 4.87s

The existing OpInfo suite is unchanged, and it could not have caught this.
sample_inputs_roll uses only nonnegative dims, and its one large shift sample
rolls a (5, 5, 5) tensor by 10000, which is a multiple of 5 and therefore the
identity. Neither roll entry in ops_test_data.py carries a skip or an xfail,
so there was nothing to remove.

Also checked by hand and passing: several dims in one call, rank 3 with
dims=-2, a shift of zero, a shift that is an exact multiple of the length, a
dimension of length zero, and a dynamic dimension whose length is only known at
run time, which is the case where the Mod node has to survive into the graph
instead of folding away.

One case worth flagging for review. Rolling a dimension whose length is zero now
evaluates Mod with a zero divisor, which the ONNX spec leaves undefined.
onnxruntime returns the dividend, so both slices come back empty and Concat
returns the input unchanged, which is what torch does. If you would rather not
depend on that, the alternative is to guard the modulo when the length is
statically zero. Separately, torch.roll(torch.zeros(2, 0), shifts=3) with no
dims still fails at Reshape with Invalid position of 0. That is unrelated to
this change and fails the same way before it.

Environment: Python 3.10, torch 2.9.1, onnx 1.22.0, onnxruntime 1.23.2, macOS
arm64.