From e8a968f748063dca5739b28d656a5249100f7b6b Mon Sep 17 00:00:00 2001
From: Ruilin Gao <229880965+ruiling-smartbear@users.noreply.github.com>
Date: Tue, 15 Sep 2026 13:19:23 +0000
Subject: [PATCH] [FIX][Relax] Guard unknown tensor shapes in VMShapeLower

Only apply the static-shape shortcut when a ShapeExpr exists. Keep unknown tensor shapes on the existing runtime-check path and cover bundled symbolic parameters.
---
 src/relax/backend/vm/vm_shape_lower.cc        |  2 +-
 .../test_backend_transform_shape_lower.py     | 71 +++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 09053a1df97f..279d5c1794d1 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -820,7 +820,7 @@ class VMShapeLowerMutator
                   const ffi::String& err_ctx, std::vector<MatchShapeTodoItem>* match_todos) final {
     // emit runtime check of shape
     auto* shape_expr = op->shape.as<ShapeExprNode>();
-    if (dynamic_only &&
+    if (dynamic_only && shape_expr != nullptr &&
         std::all_of(shape_expr->values.begin(), shape_expr->values.end(),
                     [](const PrimExpr& e) { return e->IsInstance<IntImmNode>(); })) {
       // if we only check dynamic shapes, and the shape is static, we can skip.
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py b/tests/python/relax/test_backend_transform_shape_lower.py
index 746edf9ace7a..8868bb96c96b 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -727,6 +727,77 @@ def main(
     assert_structural_equal(after, expected)
 
 
+def test_check_weights_with_unknown_shape():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(x: R.Tensor((16,), "float32"), params: R.Tuple(R.Tensor(dtype="float32", ndim=1))):
+            R.func_attr({"relax.force_pure": True, "num_input": 1})
+            return params
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor((16,), "float32"), params: R.Tuple(R.Tensor(dtype="float32", ndim=1))):
+            R.func_attr({"relax.force_pure": True, "num_input": 1})
+            shape_heap: R.Any = R.null_value()
+            _: R.Tuple = R.call_packed(
+                "vm.builtin.check_tensor_info",
+                x,
+                R.prim_value(1),
+                R.dtype("float32"),
+                R.str(""),
+                ty_args=(R.Tuple,),
+            )
+            _1: R.Tuple = R.call_packed(
+                "vm.builtin.check_tuple_info",
+                params,
+                R.prim_value(1),
+                R.str(""),
+                ty_args=(R.Tuple,),
+            )
+            gv: R.Tensor(dtype="float32", ndim=1) = params[0]
+            _2: R.Tuple = R.call_packed(
+                "vm.builtin.check_tensor_info",
+                gv,
+                R.prim_value(1),
+                R.dtype("float32"),
+                R.str(""),
+                ty_args=(R.Tuple,),
+            )
+            _3: R.Tuple = R.call_packed(
+                "vm.builtin.match_shape",
+                x,
+                shape_heap,
+                R.prim_value(1),
+                MatchShapeCode.ASSERT_EQUAL_TO_IMM,
+                R.prim_value(16),
+                R.str(""),
+                ty_args=(R.Tuple,),
+            )
+            return params
+
+    after = relax.transform.VMShapeLower(emit_err_ctx=False)(Before)
+    assert_structural_equal(after, Expected)
+
+
+def test_lower_bundled_symbolic_shape():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((32,), "float32"),
+            extent: R.Prim("int64"),
+            weight: R.Tensor(("extent",), "float32"),
+        ):
+            R.func_attr({"num_input": 1, "relax.force_pure": True})
+            return R.add(x, weight)
+
+    bundled = relax.transform.BundleModelParams()(Before)
+    lowered = relax.transform.VMShapeLower()(bundled)
+    assert isinstance(lowered, tvm.IRModule)
+
+
 def test_check_weights_with_dynamic_shape():
     MS = MatchShapeCode