[torchlib] Reimplement as_strided without an ONNX loop

[Copilot]

aten_as_strided previously lowered through a private ONNX function that built gather indices with sequence operations and a loop-like construction that was difficult to fold. This change computes the same flat-storage indices without emitting ONNX Loop or Scan.

Implementation

For each output position, the storage index is:

storage_offset + Σ_d i_d * stride[d]

The input is flattened and gathered with those indices:

• Small static views: compute the full index tensor with NumPy and emit one constant.
• Dynamic or large static views: build indices with unrolled Range/Mul/Unsqueeze/Add operations. Static constants are limited to 512 * 512 elements, matching the optimizer's default folded-output size limit.
• Normalize storage_offset=None to zero.
• Keep SymInt values as INT64 without redundant casts; this export path guarantees INT64 SymInts.

ONNX exposes logical tensor values rather than backing storage. Therefore, flattening is correct only when logical row-major order matches storage order; non-contiguous input views remain unsupported, consistent with the previous lowering.

Other changes

• Remove the unused private _aten_as_strided_onnx function and its type-constraint-test skip.
• Use Sequence[INT64] for the size and stride parameters of aten_as_strided_copy, aten_as_strided_scatter, aten_empty_strided, and aten_new_empty_strided.
• Rename the scalar reshape target from empty_shape to scalar_shape and hoist the shared trailing comment.
• Add focused coverage proving small static views retain the constant path, large static views avoid a full-size Constant/initializer after optimization, and static/dynamic results match PyTorch.

Validation

• python -m pytest -q tests/function_libs/torch_lib/e2e_ops_tests.py -k "aten_as_strided" — 7 passed
• python -m pytest -q tests/function_libs/torch_lib/ops_test.py -k "new_empty_strided" — 5 passed, 1 skipped
• python -m pytest -q onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py — 43 passed
• lintrunner -a on all three PR-changed files — no lint issues

[Justin Chu (justinchuby)]

Copilot the inputs marked as SymInt can be dynamic. Can you update the logic to account for this situation?

[Copilot]

Copilot the inputs marked as SymInt can be dynamic. Can you update the logic to account for this situation?

Added a dynamic path in 93d3247: when any of size/stride/storage_offset is a runtime value, the indices are built with ONNX ops (Range/Mul/Unsqueeze/Add) instead of a folded constant. The per-dimension contributions are still unrolled at trace time (no loop/scan) since the rank is static, and SymInt values are cast to INT64 for consistent arithmetic. The static fast path (all concrete) still folds to a single constant Gather. Verified both paths against torch.as_strided for multi-dim, non-zero storage_offset, scalar/empty-size, mixed static/dynamic, and int32/int64 SymInt inputs.

[Justin Chu (justinchuby)]

Copilot can you add good coverage on the new logic in tests/function_libs/torch_lib/e2e_ops_tests.py?

[Justin Chu (justinchuby)]

Copilot also add more detailed comment to explain what each line of the implementation is doing in relation to the torch op, as it is a more complicated op, for readability

[Copilot]

Copilot can you add good coverage on the new logic in tests/function_libs/torch_lib/e2e_ops_tests.py?

Added e2e coverage in tests/function_libs/torch_lib/e2e_ops_tests.py exercising both paths: static (multi-dim with non-zero storage_offset, single dim, overlapping strides, scalar/empty size) and dynamic (size derived from the input shape, with and without storage_offset). Writing the dynamic tests surfaced a bug where a default storage_offset=None reached the dynamic path and produced an invalid Reshape; fixed by normalizing None to 0. All tests pass against torch.as_strided.

[Copilot]

Copilot also add more detailed comment to explain what each line of the implementation is doing in relation to the torch op, as it is a more complicated op, for readability

Added detailed line-by-line comments to aten_as_strided, explaining the storage-index math in relation to torch.as_strided and annotating each step of both the static and dynamic paths.

[Copilot]

Pull request overview

This PR refactors the torchlib lowering for aten::as_strided to avoid generating an ONNX Loop/Scan when constructing gather indices, enabling simpler downstream constant-folding. It introduces a static fast path that precomputes the index tensor with NumPy when all size/stride/storage_offset values are known at trace time, and a dynamic path that builds the same index math using ONNX ops without loops.

Changes:

• Reimplemented aten_as_strided in ops/core.py as a Reshape([-1]) + Gather with (1) a NumPy-constant index fast path and (2) an ONNX-op dynamic index path (no Loop/Scan).
• Removed the now-unused private _aten_as_strided_onnx lowering and unblocked type-constraint deduction by removing it from the “skip loop/scan” list.
• Added new E2E tests covering several as_strided scenarios (static and dynamic shapes/offsets).

Reviewed changes

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

File	Description
tests/function_libs/torch_lib/e2e_ops_tests.py	Adds E2E export coverage for torch.as_strided across static and dynamic cases.
onnxscript/function_libs/torch_lib/ops/core.py	Replaces loop-based index construction with static NumPy-constant and dynamic ONNX-op paths.
onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py	Removes _aten_as_strided_onnx from the loop/scan skip list since it no longer exists.

[codecov]

Codecov Report

❌ Patch coverage is 53.57143% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.70%. Comparing base (9b6d4cc) to head (4d70d2b).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	53.57%	12 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2928      +/-   ##
==========================================
+ Coverage   72.66%   72.70%   +0.03%     
==========================================
  Files         265      265              
  Lines       32297    32298       +1     
  Branches     3056     3059       +3     
==========================================
+ Hits        23469    23481      +12     
+ Misses       7791     7779      -12     
- Partials     1037     1038       +1     

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

[Ti-Tai Wang (titaiwangms)]

Multi-reviewer summary (readability, code, critical, deep-semantic, integration)

Genuine improvement — replacing the un-foldable Loop with a single flatten + Gather is the right call. The index math (storage_offset + Σ_d i_d·stride[d]) was verified correct against torch.as_strided, including empty-size→scalar and overlapping-stride cases, on both the static NumPy path and the unrolled dynamic ONNX path. The private _aten_as_strided_onnx is cleanly removed and the deduce_type_constraints skip-list updated.

Major

• Dynamic path assumes SymInt scalars are already INT64. zero/one are INT64 constants, but size[dim], stride[dim], and storage_offset are only Reshaped, not cast. An INT32 scalar from some export path would cause Range/Mul/Add type mismatches. Suggest an explicit Cast(..., to=INT64.dtype) on the dynamic size/stride/offset scalars.
• Non-contiguous storage semantics. Reshape(self, [-1]) yields the logical row-major order, not torch's underlying storage order, so a non-contiguous self (e.g. a transposed view) would diverge. This matches the prior implementation's assumption and is normally guaranteed by dynamo decomposition, so it is a documented limitation rather than a regression — worth a one-line comment noting the contiguity assumption.

Minor

• Static path can materialize huge constants. For large/unfold-like views the static path emits an int64 index tensor of shape size, which can bloat the model / blow up exporter memory. Consider a product-size threshold that falls back to the dynamic path.
• Sibling consistency (out of scope for this PR). aten_as_strided_copy, aten_as_strided_scatter, aten_empty_strided, and aten_new_empty_strided still use the old INT64 / Sequence[int] hints rather than the new Sequence[INT64]. Possible follow-up.

Readability

• empty_shape reads as "shape of an empty tensor" but is actually the reshape-to-scalar target; rename to scalar_shape.
• The identical trailing comment appears in both branches — could be hoisted once before the if/else.

The architectural comment block explaining the storage-offset formula is excellent documentation.

[Justin Chu (justinchuby)]

Thanks. SymInts are always int64. Others I will fix

[Bas Aarts (bas-aarts)]

awesome to see this PR. I've been using a similar custom implementation for a while now.

[Justin Chu (justinchuby)]

Addressed in 4d70d2b:

• Documented that flattening relies on logical row-major order matching storage order; non-contiguous input views remain unsupported, as in the previous lowering.
• Bounded the static NumPy index path at 512 * 512 elements (matching the optimizer's folded-output limit) and fall back to the existing unrolled ONNX construction above it. Added tests proving a small static shape emits the constant index tensor and a 512 x 513 static view avoids a full-size Constant/initializer while matching PyTorch.
• Updated size/stride to Sequence[INT64] for aten_as_strided_copy, aten_as_strided_scatter, aten_empty_strided, and aten_new_empty_strided; aten_new_empty_strided now merges the sequence before ConstantOfShape.
• Renamed empty_shape to scalar_shape and hoisted the shared trailing Gather explanation.
• Intentionally did not add Cast operations: SymInts are guaranteed INT64 in this export path, so casts would be redundant.

Validation completed before commit:

• python -m pytest -q tests/function_libs/torch_lib/e2e_ops_tests.py -k "aten_as_strided" — 7 passed
• python -m pytest -q tests/function_libs/torch_lib/ops_test.py -k "new_empty_strided" — 5 passed, 1 skipped
• python -m pytest -q onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py — 43 passed
• lintrunner -a on all three PR-changed files — no lint issues

[Copilot]

🟡 Changes recommended

A newly added test compares an IR shape object directly to a Python list, which is likely to fail despite matching dimensions.

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

Review details

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