Emit scalar Range bounds in aten_unfold

[om singhal (Om-singhaI)]

aten_unfold reads the size of the unfolded dimension like this:

dim_size = op.Gather(input_shape, op.Constant(value_ints=[dimension]))
window_starts = op.Range(0, op.Sub(dim_size, size - 1), step)

A rank 1 index makes Gather return int64[1], not a scalar. Sub preserves that rank, so a rank 1 value reaches the limit input of Range. The ONNX spec requires all three Range inputs to be rank 0, and onnxruntime rejects the model at session creation:

Node (node_Range_6) Op (Range) [ShapeInferenceError] Input to 'Range' op should be scalars (Tensor with only one element and shape empty)

Scope

This is narrower than it first looks. With x.unfold(1, 2, 1) on a (3, 4) input:

export	ops in graph	onnxruntime
static dim, optimize=True (default)	['Gather']	runs, correct
static dim, optimize=False	full chain, Range survives	rejected at session creation
dynamic unfold dim	full chain, Range survives	runs

On the default path the constant folder collapses Shape/Gather/Sub/Range into a single index constant, so no Range node survives and those exports are fine and numerically correct. The bad node is reachable with optimize=False, and on a dynamic unfold dimension Range survives with a rank 1 limit but onnxruntime cannot fold its inputs, so the scalar check never runs and the spec violation is silently tolerated.

So: a spec compliance fix on an unoptimized path, not a broken operator.

The fix

# Range requires rank 0 (scalar) inputs, so squeeze the [1] shaped
# Gather result down to a scalar.
dim_size = op.Squeeze(
    op.Gather(input_shape, op.Constant(value_ints=[dimension])),
    op.Constant(value_ints=[0]),
)

Why not just value_int=dimension

That one token change also gives Range a scalar and it does fix onnxruntime, but it quietly defeats the Shape/Gather constant folding. The rule in onnxscript/optimizer/_constant_folding.py bails out on any index whose ndim is not 1:

if indices_numpy_value.ndim != 1:
    return None

Measured on a dynamic batch dimension with a static unfold dimension, default optimize=True:

main                    ops: ['Gather']
value_int variant       ops: ['Shape', 'Gather', 'Sub', 'Range', 'Unsqueeze', 'Add', 'Gather']
Squeeze variant         ops: ['Gather']

Squeeze is registered as a shape value propagator in that same file, so the symbolic shape value survives the squeeze and the folded path keeps working. The seven node chain would otherwise be left in the graph at runtime.

Testing

Added test_unfold_emits_scalar_range_bounds to tests/function_libs/torch_lib/e2e_ops_tests.py, following the optimize=False pattern already used by the tests around it.

With only the core.py change reverted, the new test fails with the error quoted above. With the fix it passes. Rest of e2e_ops_tests.py and the unfold cases in ops_test.py are unchanged.

Same class of fix as #2943 (rank 1 where a scalar was required, same file) and #3006 (invalid ONNX graph out of a torchlib lowering).

[Justin Chu (justinchuby)]

Thanks!

[Copilot]

Pull request overview

This PR fixes an ONNX spec-compliance issue in the aten_unfold lowering where the limit input to Range could become a rank-1 tensor (int64[1]) due to Gather returning [1]-shaped output for a rank-1 index. The change ensures Range receives rank-0 scalar inputs (as required by the ONNX spec), preventing onnxruntime from rejecting unoptimized exports.

Changes:

• Update aten_unfold to Squeeze the [1]-shaped Gather result into a scalar before feeding it to Range.
• Add an end-to-end regression test that exports with optimize=False to ensure the scalar-Range-bounds behavior is exercised.

Reviewed changes

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

File	Description
onnxscript/function_libs/torch_lib/ops/core.py	Squeezes the gathered dimension size to a scalar to satisfy Range’s scalar-input requirement.
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds an optimize=False export regression test for unfold to catch invalid Range inputs.

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.63%. Comparing base (e1fe520) to head (1bd9cbf).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3013   +/-   ##
=======================================
  Coverage   72.63%   72.63%           
=======================================
  Files         265      265           
  Lines       32218    32218           
  Branches     3044     3044           
=======================================
  Hits        23403    23403           
  Misses       7781     7781           
  Partials     1034     1034           

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

[om singhal (Om-singhaI)]

Thanks! Those two red checks are the py311 torch nightly jobs, they're red on main and on every other open PR at the moment so there's nothing to chase there.