Fix aten_roll for negative dims and shifts past the dimension length by Om-singhaI · Pull Request #3024 · microsoft/onnxscript

om singhal (Om-singhaI) · 2026-08-30T03:30:55Z

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.

`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: ```python # 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`: ```python 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: ```python 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. Environment: Python 3.10, torch 2.9.1, onnx 1.22.0, onnxruntime 1.23.2, macOS arm64.

Copilot

Pull request overview

This PR fixes the aten::roll ONNX lowering to match torch.roll semantics by (1) normalizing negative dims and (2) making shifts truly circular via modulo, which also resolves an onnxruntime model-load failure for dims=-1.

Changes:

Normalize negative dims for aten_roll and adjust negative-dim normalization for the complex variant to account for the trailing complex axis.
Replace “Slice clamping” behavior with explicit modulo-based wrap for both per-dim roll and no-dim (flattened) roll paths.
Add E2E export/runtime tests covering negative dims and shifts that exceed the dimension length (including the dims=() path).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
`onnxscript/function_libs/torch_lib/ops/core.py`	Fixes roll lowering by normalizing negative dims and computing circular wrap via modulo for per-dim and no-dim paths.
`tests/function_libs/torch_lib/e2e_ops_tests.py`	Adds regression tests covering negative dims and large positive/negative shifts for real and complex roll exports.

Suppressed comments (1)

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

slice_length uses op.Mod(..., dim_length). If the rolled dimension has length 0, this becomes Mod(x, 0), which is undefined in ONNX and can lead to runtime errors or inconsistent behavior across runtimes. Guarding the modulo (treating length-0 as slice_length=0) keeps roll semantics well-defined for zero-length dimensions.

    dim_length = op.Shape(self, start=dim, end=dim + 1)
    # roll is circular, so the shift is taken modulo the length of the dimension
    slice_length = op.Mod(op.Constant(value_ints=[-shift]), dim_length)

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

codecov · 2026-08-31T17:43:11Z

Codecov Report

❌ Patch coverage is 73.33333% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.64%. Comparing base (3ba2bf7) to head (483b8a9).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	73.33%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3024   +/-   ##
=======================================
  Coverage   72.64%   72.64%           
=======================================
  Files         265      265           
  Lines       32251    32263   +12     
  Branches     3050     3050           
=======================================
+ Hits        23429    23439   +10     
- Misses       7786     7787    +1     
- Partials     1036     1037    +1

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

justinchuby asked for an assertion rather than a guard in the graph. Both helpers now assert that the length they divide by is not zero, at the two modulo sites, in the same shape as the other preconditions in this file. For the assertion to be a precondition rather than a trap, the empty tensor has to be caught before it gets that far. The existing early return only looked at the first dimension, so a shape like (2, 0) still reached the modulo. It now tests every dimension, in aten_roll and in aten_roll_complex alike, and an empty tensor exports to a single Identity, which is what torch returns for it. That also settles the loose end in the description of this pull request. On 35f22c4, torch.roll(torch.zeros(2, 0), shifts=3) with no dims exported a graph that failed at Reshape with "Invalid position of 0" under onnxruntime. It now returns the input unchanged. Nine zero length cases across both helpers and both variants, real and complex, now match torch. Four of them previously depended on onnxruntime returning the dividend for Mod by zero, which the ONNX spec leaves undefined, and one failed outright. The graph for a nonempty tensor is untouched, still 8 nodes for the per dim path and 11 for the flattened path. A dimension whose length is only zero at run time is outside what a trace time assertion can see. The Mod node survives into the graph for a dynamic dimension, as before, and that case still leans on the runtime. Added test_roll_empty_tensor_is_an_identity and its complex counterpart. They assert the exported graph carries no Mod at all and then compare against torch. On 35f22c4 three of the five fail. pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k roll 12 passed, 100 deselected pytest tests/function_libs/torch_lib/ops_test.py -k roll 6 passed, 2 skipped, 1844 deselected

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

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

_aten_roll_shift_and_dim_onnx computes slice_length = Mod(-shift, dim_length). When dim_length is only known at runtime (dynamic shapes), it may be 0 for empty tensors, which again makes Mod(x, 0) undefined. Clamping the divisor to at least 1 keeps the result well-defined while still producing empty slices (and thus an identity roll) when the actual dimension length is 0.

    dim_length = op.Shape(self, start=dim, end=dim + 1)
    # roll is circular, so the shift is taken modulo the length of the dimension
    slice_length = op.Mod(op.Constant(value_ints=[-shift]), dim_length)

+    total_length = op.Shape(self_flatten)
+    # Compute slice length. roll is circular, so the shift is taken modulo the number
+    # of elements. For [A,B,C,D], if shift is 1, slice_length = 3, means move [A,B,C]
+    # to the end. The effect equals to move [D] to the beginning.
+    slice_length = op.Mod(op.Constant(value_ints=[-shift]), total_length)


om singhal (Om-singhaI) · 2026-09-01T18:31:13Z

+    prefix = op.Slice(self_flatten, slice_length, total_length)
    # Concat first+second together, e.g. [D,A,B,C]
    result = op.Concat(prefix, suffix, axis=0)
    return op.Reshape(result, op.Shape(self))


Added allowzero. It costs no nodes and stops a zero in the target shape being read as copy the input dimension, which is what gave the Invalid position of 0.

Left the clamp out, that is the runtime guard ruled out above, and a divisor that only turns out to be zero at run time is past what the assert can see anyway.

…zero A dimension that is only zero at run time cannot be caught by the assert above, and Reshape reads a zero in the target shape as "copy the input dimension at this index" unless allowzero is set. Against the flattened tensor that produces "Invalid position of 0" rather than an empty result. The attribute costs no nodes and changes nothing for a shape with no zeros in it.

Copilot

🟡 Changes recommended

The new empty-tensor tests iterate over onnx_program.model.graph (non-iterable) instead of onnx_program.model.graph.node, which will raise at runtime and fail CI.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

tests/function_libs/torch_lib/e2e_ops_tests.py:996

onnx_program.model.graph is a GraphProto and isn’t iterable; this list comprehension will raise at runtime. Iterate over onnx_program.model.graph.node instead.

        self.assertNotIn("Mod", [node.op_type for node in onnx_program.model.graph])

Files reviewed: 2/2 changed files
Comments generated: 1
Review effort level: Lite

+        onnx_program = torch.onnx.export(
+            RollModel(), (torch.zeros(shape),), dynamo=True, optimize=False
+        )
+        self.assertNotIn("Mod", [node.op_type for node in onnx_program.model.graph])


github-project-automation Bot added this to ONNX Script Review Board Aug 30, 2026

github-project-automation Bot moved this to Todo in ONNX Script Review Board Aug 30, 2026

Justin Chu (justinchuby) requested a lite review from Copilot August 31, 2026 17:36

Justin Chu (justinchuby) added the module: torchlib Related to the torch/aten function lib in development label Aug 31, 2026

Justin Chu (justinchuby) self-assigned this Aug 31, 2026

Copilot started reviewing on behalf of Justin Chu (justinchuby) August 31, 2026 17:37 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

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

Justin Chu (justinchuby) requested a lite review from Copilot September 1, 2026 01:49

Copilot started reviewing on behalf of Justin Chu (justinchuby) September 1, 2026 01:49 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Justin Chu (justinchuby) requested a lite review from Copilot September 3, 2026 04:25

Copilot started reviewing on behalf of Justin Chu (justinchuby) September 3, 2026 04:26 View session

Copilot AI reviewed Sep 3, 2026

View reviewed changes

Comment thread tests/function_libs/torch_lib/e2e_ops_tests.py

onnx_program = torch.onnx.export(

RollModel(), (torch.zeros(shape),), dynamo=True, optimize=False

)

self.assertNotIn("Mod", [node.op_type for node in onnx_program.model.graph])

Navigation Menu

Uh oh!

Fix aten_roll for negative dims and shifts past the dimension length - #3024

Fix aten_roll for negative dims and shifts past the dimension length#3024
om singhal (Om-singhaI) wants to merge 3 commits into
microsoft:mainfrom
Om-singhaI:fix/roll-negative-dim-and-shift-wrap

om singhal (Om-singhaI) commented Aug 30, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

codecov Bot commented Aug 31, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

om singhal (Om-singhaI) Sep 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

om singhal (Om-singhaI) commented Aug 30, 2026

A negative dim emits a model onnxruntime will not load

A shift past the dimension length is silently wrong

The fix

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

codecov Bot commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

om singhal (Om-singhaI) Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Aug 31, 2026 •

edited

Loading