From bf23939b3ba1b59324a044c67bc0bfac6fe2dd21 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Wed, 12 Aug 2026 03:02:49 -0400 Subject: [PATCH] [FIX][RELAX][Torch] Materialize runtime scalar shape values --- .../torch/base_fx_graph_translator.py | 41 ++++++++++++++--- .../torch/exported_program_translator.py | 3 ++ .../test_frontend_from_exported_program.py | 45 +++++++++++++++++++ 3 files changed, 84 insertions(+), 5 deletions(-) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 85a0d956ca1d..4949ec125916 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2498,6 +2498,13 @@ def _fill(self, node: fx.Node) -> relax.Var: value = args[1] if isinstance(args[1], relax.Expr) else relax.const(args[1], dtype) return self.block_builder.emit(relax.op.full(x.ty.shape, value, dtype)) + def _prim_value_to_scalar_tensor(self, value: relax.Expr, dtype: str) -> relax.Var: + """Materialize an integer primitive value as a rank-zero tensor.""" + value_tensor = self.block_builder.emit(relax.op.shape_to_tensor(relax.ShapeExpr([value]))) + if dtype != "int64": + value_tensor = self.block_builder.emit(relax.op.astype(value_tensor, dtype)) + return self.block_builder.emit(relax.op.squeeze(value_tensor, axis=[0])) + def _inplace_fill(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] @@ -2515,7 +2522,12 @@ def _full(self, node: fx.Node) -> relax.Var: dtype = self._convert_data_type( node.kwargs.get("dtype", torch.get_default_dtype()), self.env ) - value = args[1] if isinstance(args[1], relax.expr.Constant) else relax.const(args[1], dtype) + if isinstance(getattr(args[1], "ty", None), PrimType): + value = self._prim_value_to_scalar_tensor(args[1], dtype) + elif isinstance(args[1], relax.Expr): + value = args[1] + else: + value = relax.const(args[1], dtype) return self.block_builder.emit( relax.op.full( size, @@ -2525,11 +2537,17 @@ def _full(self, node: fx.Node) -> relax.Var: ) def _full_like(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - value = node.args[1] - fill_value = relax.const(value) + args = self.retrieve_args(node) + x = args[0] + value = args[1] + x_dtype = str(x.ty.dtype) + if isinstance(getattr(value, "ty", None), PrimType): + fill_value = self._prim_value_to_scalar_tensor(value, x_dtype) + elif isinstance(value, relax.Expr): + fill_value = value + else: + fill_value = relax.const(value) - x_dtype = x.ty.dtype.dtype fill_dtype = None if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)): if not ("float" in x_dtype or "bfloat16" in x_dtype): @@ -2800,6 +2818,19 @@ def _getitem(self, node: fx.Node) -> relax.Var: def _item(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] + shape = self.shape_of(x) + if shape is not None and len(shape) == 0: + dtype = str(x.ty.dtype) + if dtype in ("int32", "int64"): + scalar = x + if dtype != "int64": + scalar = self.block_builder.emit(relax.op.astype(scalar, "int64")) + scalar = self.block_builder.emit(relax.op.reshape(scalar, [1])) + shape_value = self.block_builder.emit(relax.op.tensor_to_shape(scalar)) + dim = tirx.Var(f"{node.name}_dim", "int64") + self.block_builder.match_cast(shape_value, relax.ShapeType([dim])) + return dim + return x return self.block_builder.emit(relax.op.take(x, relax.const(0, "int64"), axis=0)) def _sym_size_int(self, node: fx.Node) -> relax.Expr: diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index ea604c21cc84..4e614da6a8d8 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1792,6 +1792,7 @@ def create_convert_map( "triu.default": self._tril_triu(relax.op.triu), "trunc.default": self._unary_op(relax.op.trunc), # binary + "add": self._binary_op(relax.op.add, operator.add), "add.Tensor": self._binary_op(relax.op.add, operator.add), "add.Scalar": self._binary_op(relax.op.add, operator.add), "add_.Tensor": self._binary_op(relax.op.add, operator.add), @@ -1847,6 +1848,7 @@ def create_convert_map( "pow.Scalar": self._binary_op(relax.op.power, operator.pow), "pow.Tensor_Scalar": self._pow, "pow.Tensor_Tensor": self._binary_op(relax.op.power, operator.pow), + "sub": self._binary_op(relax.op.subtract, operator.sub), "sub.Tensor": self._binary_op(relax.op.subtract, operator.sub), "sub.Scalar": self._binary_op(relax.op.subtract, operator.sub), "__and__.Tensor": self._binary_op(relax.op.bitwise_and, operator.and_), @@ -2038,6 +2040,7 @@ def create_convert_map( "le": self._symbolic_comparison, "gt": self._symbolic_comparison, "lt": self._symbolic_comparison, + "eq": self._symbolic_comparison, # higher-order ops "cond": self._cond, } diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 3799b8ed95f5..6d412430be74 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5507,6 +5507,51 @@ def main(x: R.Tensor((2, 8, 4), dtype="float32")) -> R.Tuple( verify_model(SliceStaticModel(), example_args_static, {}, ExpectedStatic) +def test_dynamic_scalar_item_in_shape_operations(): + class DynamicShapeOps(torch.nn.Module): + def forward(self, x): + lengths = torch.full( + (x.shape[0],), + x.shape[1], + device=x.device, + dtype=torch.int64, + ) + max_len = lengths.max().item() + positions = torch.arange(max_len, device=x.device) + mask = positions.unsqueeze(0).expand(x.shape[0], -1) == lengths.unsqueeze(1) + filled = torch.full( + (x.shape[0], x.shape[1]), + x.shape[1], + device=x.device, + dtype=torch.int64, + ) + shifted = torch.arange(x.shape[1] + 1, device=x.device) + shortened = torch.full_like(lengths, x.shape[1] - 1) + return mask, filled, shifted, shortened + + example_args = (torch.randn(1, 4, 3, dtype=torch.float32),) + tokens = torch.export.Dim("tokens", min=1, max=8) + exported_program = export( + DynamicShapeOps(), + args=example_args, + dynamic_shapes={"x": {1: tokens}}, + ) + mod = from_exported_program(exported_program) + + script = mod.script() + assert "R.tensor_to_shape" in script + assert "R.shape_to_tensor" in script + + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + for token_count in (4, 6): + torch_input = torch.randn(1, token_count, 3) + expected = DynamicShapeOps()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + def test_split(): class Chunk(Module): def forward(self, input):