Skip to content

[Fix][Relax] Preserve symbolic shapes during constant folding - #20431

Open
byy723 wants to merge 4 commits into
apache:mainfrom
byy723:repro/issue-20260
Open

byy723 wants to merge 4 commits into
apache:mainfrom
byy723:repro/issue-20260

Conversation

@byy723

@byy723 byy723 commented Sep 24, 2026

Copy link
Copy Markdown

Motivation

For valid symbolic shapes, the shape_to_tensor fallback in FoldConstant can encounter a PrimExpr that is not an IntImmNode. Dereferencing it as an integer constant is unsafe. In the valid shape_of -> shape_to_tensor -> tensor_to_shape -> reshape path, reshape legalization also assumes a shape Var is bound directly to a ShapeExpr, but it may be bound to a call.

Related to #20260.

Changes

  • Fold shape_to_tensor through the C++ fallback only if each dimension is an int64 IntImmNode with a value representable as int64_t; otherwise retain the call.
  • Preserve reshape when its target shape Var is not bound to a ShapeExpr, rather than asserting during legalization.
  • Add regression tests for symbolic and mixed shapes, static folding through the fallback, and a valid dynamic reshape. The fallback tests temporarily remove and restore FLegalize to exercise the C++ path.

Validation

  • The complete tests/python/relax/test_transform_fold_constant.py file passes: 23 tests on the rebased source.
  • The original valid dynamic-reshape reproducer completes FoldConstant without an exception.
  • Both the default and cpu_generic LLVM/CPU pipelines build and execute with input lengths 3 and 7; outputs match the inputs.
  • Pre-commit checks pass on the three changed files.

Scope

Passing the Tensor result of shape_to_tensor directly to reshape without tensor_to_shape remains outside this change. In the BlockBuilder reproducer, that input is rejected before FoldConstant because reshape expects a Shape-typed operand. This PR does not relax that contract.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant