diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 27d30e377f..02fde7b5c9 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8791,8 +8791,9 @@ def aten_roll(self: TTensor, shifts: Sequence[int], dims: Sequence[int] = ()) -> if self_rank == 0: return op.Identity(self) elif 0 in self.shape: - # A tensor with no elements rolls to itself, which is what torch returns. It is - # also what keeps a zero length out of the modulo in the helpers below. + # A tensor with no elements rolls to itself, which is what torch returns. It also + # keeps a zero length out of the modulo in the helpers below, but only when the + # length is known to be zero at export time. See the limitation noted there. return op.Identity(self) # NOTE: In pytorch, default value of dims is an empty list. @@ -8858,6 +8859,12 @@ def aten_roll_complex( def _aten_roll_shift_no_dim_onnx(self: TTensor, shift: int) -> TTensor: # The element count is the divisor of the Mod below, and Mod by zero is undefined in # ONNX. Both callers return a tensor with no elements unchanged before reaching here. + # Known limitation: that early return and this assert only cover a length known to be + # zero at export time. A dynamic dimension that is zero at run time still reaches Mod + # with a zero divisor, which ONNX leaves undefined. Runtimes differ: the onnxruntime CPU + # kernel rejects an integer divisor of zero, its CUDA kernel has no such check, and + # onnxruntime before 1.25 returned the dividend, which the Slice calls below clamp back + # into range so the result still looks correct. assert self.shape is None or 0 not in self.shape, ( "the element count must not be zero because Mod by zero is undefined" ) @@ -8887,6 +8894,12 @@ def _aten_roll_shift_and_dim_onnx(self: TTensor, shift: int, dim: int) -> TTenso # The length of that dimension is the divisor of the Mod below, and Mod by zero is # undefined in ONNX. Both callers return a tensor with no elements unchanged before # reaching here. + # Known limitation: that early return and this assert only cover a length known to be + # zero at export time. A dynamic dimension that is zero at run time still reaches Mod + # with a zero divisor, which ONNX leaves undefined. Runtimes differ: the onnxruntime CPU + # kernel rejects an integer divisor of zero, its CUDA kernel has no such check, and + # onnxruntime before 1.25 returned the dividend, which the Slice calls below clamp back + # into range so the result still looks correct. assert self.shape is None or self.shape[dim] != 0, ( "the dimension length must not be zero because Mod by zero is undefined" )