[FIX][Relax][ONNX] Keep the static shape of a rank-0 Shape input

[adityasingh2400]

Shape._impl_v13 in the Relax ONNX frontend chose between the static shape and a runtime shape_of with a truthiness test:

if not data_info.shape:
    return bb.normalize(relax.op.shape_of(inputs[0]))
return data_info.shape

A rank-0 tensor has a defined but empty shape, and an empty ShapeExpr is falsy, so a scalar input took the runtime path that is meant only for a tensor whose shape is genuinely unknown. TensorType.shape is None in that case and R.shape([]) for a rank-0 tensor, so the two are distinguishable, and only the first should reach shape_of.

The value handed back was then a normalized R.shape_of call rather than a ShapeExpr, so every downstream converter that matches on relax.ShapeExpr lost its static path. Slice is the visible case from the report: importing Shape followed by Slice on a scalar input raises

Error converting operator Slice, with inputs: [R.shape_of(X), ...]
ValueError: Slice requires a statically known input rank.

because _get_known_tensor_rank cannot produce a rank for a shape value. Gather and Reshape carry the same isinstance(..., relax.ShapeExpr) match, so they are exposed to the same shape.

Fix

Compare against None, so a rank-0 input keeps the static R.shape([]) that every other rank already gets.

Shape of a scalar now folds at import time to R.shape([]), and Shape followed by Slice folds to an empty int64 tensor of shape (0,). Checked against ONNX Runtime 1.24 on the same graph: Shape of a scalar returns an empty int64 array of shape (0,), and slicing it returns the same, so the folded result matches.

Effect on an existing test

test_shape_start_end_scalar, added by #20050, pinned the runtime fallback for a rank-0 input with start=1, asserting the op chain relax.shape_of, relax.shape_to_tensor, relax.strided_slice, relax.tensor_to_shape. With the static shape preserved, that case folds to the same empty static shape, so the test now asserts the folded module and that no ops remain. ONNX Runtime returns an empty int64 array for that graph too, so the folded answer is the correct one, and the assertion change is the point of the fix rather than a workaround for it.

Testing

Two new tests in tests/python/relax/test_frontend_onnx.py:

• test_shape_scalar_input, structural equality against the expected module, pins that Shape of a rank-0 input emits R.shape([]) and not R.shape_of.
• test_slice_of_scalar_shape, the reported pattern end to end, pins that the import succeeds and yields an empty int64 tensor.

Verified fail-before and pass-after against the base ref rather than a stash, over the whole file so any collateral damage would show:

python -m pytest tests/python/relax/test_frontend_onnx.py -q   # with the fix
python -m pytest tests/python/relax/test_frontend_onnx.py -q   # at upstream/main

The failure sets differ by exactly three entries, all of them the target tests, and there are no new failures:

only in the fixed run:   (none)
only in the base run:    test_shape_scalar_input
                         test_shape_start_end_scalar
                         test_slice_of_scalar_shape

This was run on a local build configured with USE_LLVM OFF, so the tests that go through check_correctness and tvm.compile(target="llvm") fail identically in both runs with ValueError: Cannot find global function target.build.llvm. They are the same 174 entries on both sides and are unrelated to this change. The three tests above and the surrounding test_shape and test_shape_start_end cases need no codegen and were run directly, 16 passed.

Lint checked with the pinned ruff==0.12.3 from .pre-commit-config.yaml: ruff format --check reports already formatted and ruff check passes on both files.

Fixes #17770

[tlopex]

get_info still maps both an absent ONNX shape (unknown rank) and an empty shape (scalar) to []. This change would therefore fold unknown-rank values to R.shape([]). Please preserve the distinction using HasField("shape") and add an unknown-rank test.

[adityasingh2400]

You are right, and thanks for catching it. Fixed in dd1477d.

I confirmed the collapse you describe. get_info starts from shape = [] and only appends per dim, so it never returns None, and _parse_graph_input feeds that straight into _new_var. Both an absent shape field and a present-but-empty one therefore became R.Tensor(()), which means my is None check was unreachable for an unknown-rank input and such an input would have folded to R.shape([]).

get_info now returns None for the shape when tensor_type.HasField("shape") is false. That matches what TensorType already documents, shape: Optional[Expr] with ndim defaulting to -1 for the unknown-rank form, so _new_var(shape=None) produces a tensor with no static shape and the is None branch in Shape becomes reachable and correct. The call site now also warns on unknown rank, matching the existing warning for unknown dimensions.

I checked the premise against real protos rather than assuming:

input	HasField("shape")	dims
make_tensor_value_info(..., []) rank-0	True	0
make_tensor_value_info(..., None) unknown rank	False	0
make_tensor_value_info(..., [2, 3])	True	2

Added test_shape_unknown_rank_input, which asserts the input keeps shape is None and ndim == -1 rather than becoming rank-0, and that Shape still emits relax.shape_of. It also asserts the absent shape field up front so the test cannot silently stop covering the case.

One caveat on verification: I could not run the ONNX frontend tests locally, since that needs a built TVM. The proto behaviour above is measured, the rest is reasoned from get_info's single call site and the TensorType contract, so CI is the real check here. Happy to iterate if it comes back unhappy.

[adityasingh2400]

cpu/pr-head went red on dd1477d and the cause was my own test, not the frontend change. Fixed in 829262a.

The failure was a single test out of the run:

FAILED tests/python/relax/test_frontend_onnx.py::test_shape_unknown_rank_input
  - AttributeError: 'Var' object has no attribute 'struct_info'
= 1 failed, 10970 passed, 616 skipped, 57 xfailed, 6 xpassed

It reproduced identically on both test shards, and the Build stage itself passed, so this was the assertion and not the conversion.

I wrote the assertion as params[0].struct_info, which is not the API on this base. struct_info does not appear anywhere under python/tvm/relax here, since Relax vars carry a Type. So the line raised before it checked anything. It now reads params[0].ty, which exposes the same two fields:

data_ty = tvm_model["main"].params[0].ty
assert data_ty.shape is None
assert data_ty.ndim == -1

That matches how existing tests read it, for example arg.ty.dtype further down this same file and before.params[0].ty.shape in test_bind_symbolic_vars.py. It also lines up with relax.TensorType.__init__, where an omitted shape stays None and ndim defaults to -1, which is exactly what _new_var now receives for an input with no shape field.

Being clear about what I verified: I do not have a TVM build in this environment, so I could not run the test locally, and I am relying on CI for that. What I did check is that .ty is the accessor used by passing tests on this base, that TensorType leaves shape as None and ndim at -1 for this input, and that the change is limited to the three assertion lines with no formatting drift.

The frontend change from dd1477d is untouched, and arm, gpu, docker, and wasm were all green on that commit.

[adityasingh2400]

Following up on my last comment, CI has confirmed the fix and all five checks are green: arm, cpu, docker, gpu, wasm.

I said there that I had no TVM build in my environment, could not run the test locally, and was relying on CI for that. The cpu run has now finished:

= 10971 passed, 616 skipped, 57 xfailed, 6 xpassed, 630 warnings in 1631.00s =

against the previous run's 1 failed, 10970 passed, so test_shape_unknown_rank_input passes and nothing else moved. Reading the type through params[0].ty was the right call, and the frontend change from dd1477d was never involved.

@tlopex this is still marked as changes requested from your review, and I believe both points are now covered. dd1477d fixed the collapse you identified, where get_info mapped both an absent shape and an empty shape to [], and 829262a fixed my own test that was asserting through an API this base does not have. Whenever you have a moment, could you take another look?