This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new cb8fd21742 [Fix][Relax][Torch] Preserve derived exported input
dimensions (#20128)
cb8fd21742 is described below
commit cb8fd21742078184b880d43ff8fc794aafd21557
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Aug 22 11:42:15 2026 -0700
[Fix][Relax][Torch] Preserve derived exported input dimensions (#20128)
PyTorch exported programs may describe input dimensions using
expressions derived from other symbolic dimensions, such as `2 * n`. The
importer previously recognized these expressions but replaced them with
independent TIR variables when constructing the Relax function
signature. This discarded the relationship between the dimensions.
This PR:
- Preserves derived TIR expressions when constructing input tensor
shapes
- Records range constraints only for direct symbolic variables
- Accepts input shape and dtype metadata from either `tensor_meta` or
placeholder value metadata
- Reports an explicit error when input shape or dtype metadata is
unavailable
---
.../frontend/torch/exported_program_translator.py | 30 +++++++----
.../relax/test_frontend_from_exported_program.py | 61 ++++++++++++++--------
2 files changed, 59 insertions(+), 32 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index ea604c21cc..0e69074af6 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -2108,10 +2108,11 @@ class ExportedProgramImporter(BaseFXGraphImporter):
None if math.isinf(float(value_range.upper)) else
int(value_range.upper)
)
- symbol_name, _ = self._process_derived_symbol(
+ symbol_name, derived_expr =
self._process_derived_symbol(
symbol, torch_symbol_to_relax_var
)
- range_constraints[symbol_name] = (lower, upper)
+ if derived_expr is None:
+ range_constraints[symbol_name] = (lower, upper)
except (OverflowError, AttributeError, TypeError):
continue
@@ -2119,14 +2120,17 @@ class ExportedProgramImporter(BaseFXGraphImporter):
named_buffers = OrderedDict(exported_program.named_buffers())
for spec in exported_program.graph_signature.input_specs:
name_hint = spec.arg.name
+ torch_shape = None
+ torch_dtype = None
if spec.kind is
torch.export.graph_signature.InputKind.CONSTANT_TENSOR:
torch_shape =
exported_program.tensor_constants[spec.target].shape
torch_dtype =
exported_program.tensor_constants[spec.target].dtype
elif spec.kind is
torch.export.graph_signature.InputKind.USER_INPUT:
for node in
exported_program.graph.find_nodes(op="placeholder", target=spec.target):
- if node.name == name_hint and "tensor_meta" in node.meta:
- torch_shape = node.meta["tensor_meta"].shape
- torch_dtype = node.meta["tensor_meta"].dtype
+ tensor_meta = node.meta.get("tensor_meta",
node.meta.get("val"))
+ if node.name == name_hint and tensor_meta is not None:
+ torch_shape = tensor_meta.shape
+ torch_dtype = tensor_meta.dtype
break
elif spec.kind is torch.export.graph_signature.InputKind.BUFFER:
torch_shape = named_buffers[spec.target].shape
@@ -2136,19 +2140,23 @@ class ExportedProgramImporter(BaseFXGraphImporter):
torch_dtype = exported_program.state_dict[spec.target].dtype
else:
raise ValueError(f"Unsupported input kind: {spec.kind}")
+ if torch_shape is None or torch_dtype is None:
+ raise ValueError(f'Cannot determine shape and dtype for input
"{name_hint}"')
relax_shape = []
for s in torch_shape:
if isinstance(s, torch.SymInt):
sympy_node = s.node.expr if hasattr(s.node, "expr") else
s.node
- symbol_name, _ = self._process_derived_symbol(
+ symbol_name, derived_expr = self._process_derived_symbol(
sympy_node, torch_symbol_to_relax_var
)
-
- shape_var = torch_symbol_to_relax_var.setdefault(
- symbol_name, tvm.tirx.Var(symbol_name, "int64")
- )
- relax_shape.append(shape_var)
+ if derived_expr is not None:
+ relax_shape.append(derived_expr)
+ else:
+ shape_var = torch_symbol_to_relax_var.setdefault(
+ symbol_name, tvm.tirx.Var(symbol_name, "int64")
+ )
+ relax_shape.append(shape_var)
else:
relax_shape.append(s)
dtype = self._convert_data_type(torch_dtype)
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 3799b8ed95..b2aa982fdf 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -5507,6 +5507,28 @@ def test_slice_with_symbolic_end():
verify_model(SliceStaticModel(), example_args_static, {}, ExpectedStatic)
+def test_derived_input_dimension_without_exported_program_decomposition():
+ class IdentityPair(torch.nn.Module):
+ def forward(self, x, y):
+ return x, y
+
+ frames = torch.export.Dim("frames", min=1, max=8)
+ exported_program = export(
+ IdentityPair(),
+ args=(torch.randn(1, 4, 3), torch.randn(1, 8, 3)),
+ dynamic_shapes=({1: frames}, {1: 2 * frames}),
+ )
+ mod = from_exported_program(
+ exported_program,
+ keep_params_as_input=True,
+ run_ep_decomposition=False,
+ )
+
+ x_shape = mod["main"].params[0].ty.shape.values
+ y_shape = mod["main"].params[1].ty.shape.values
+ assert tvm.arith.Analyzer().can_prove_equal(y_shape[1], x_shape[1] * 2)
+
+
def test_split():
class Chunk(Module):
def forward(self, input):
@@ -8102,19 +8124,18 @@ def test_dynamic_shape_with_addition_constraints():
class Expected:
@R.function
def main(
- x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("s0___1",
4), dtype="float32")
- ) -> R.Tuple(R.Tensor(("s0 + s0___1", 4), dtype="float32")):
+ x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("1 + s0",
4), dtype="float32")
+ ) -> R.Tuple(R.Tensor(("s0 + (1 + s0)", 4), dtype="float32")):
s0 = T.int64()
- s0___1 = T.int64()
R.func_attr(
{
- "tir_var_lower_bound": {"s77": 1, "s77___1": 2},
- "tir_var_upper_bound": {"s77": 64, "s77___1": 65},
+ "tir_var_lower_bound": {"s77": 1},
+ "tir_var_upper_bound": {"s77": 64},
}
)
with R.dataflow():
- lv: R.Tensor((s0 + s0___1, 4), dtype="float32") = R.concat((x,
y), axis=0)
- gv: R.Tuple(R.Tensor((s0 + s0___1, 4), dtype="float32")) =
(lv,)
+ lv: R.Tensor((s0 + (1 + s0), 4), dtype="float32") =
R.concat((x, y), axis=0)
+ gv: R.Tuple(R.Tensor((s0 + (1 + s0), 4), dtype="float32")) =
(lv,)
R.output(gv)
return gv
@@ -8136,19 +8157,18 @@ def test_dynamic_shape_with_subtraction_constraints():
class Expected:
@R.function
def main(
- x: R.Tensor(("s0___1", 4), dtype="float32"), y: R.Tensor(("s0",
4), dtype="float32")
- ) -> R.Tuple(R.Tensor(("s0___1 + s0", 4), dtype="float32")):
- s0___1 = T.int64()
+ x: R.Tensor(("1 + s0", 4), dtype="float32"), y: R.Tensor(("s0",
4), dtype="float32")
+ ) -> R.Tuple(R.Tensor(("1 + s0 + s0", 4), dtype="float32")):
s0 = T.int64()
R.func_attr(
{
- "tir_var_lower_bound": {"s17": 0, "s17___1": 1},
- "tir_var_upper_bound": {"s17": 63, "s17___1": 64},
+ "tir_var_lower_bound": {"s17": 0},
+ "tir_var_upper_bound": {"s17": 63},
}
)
with R.dataflow():
- lv: R.Tensor((s0___1 + s0, 4), dtype="float32") = R.concat((x,
y), axis=0)
- gv: R.Tuple(R.Tensor((s0___1 + s0, 4), dtype="float32")) =
(lv,)
+ lv: R.Tensor((1 + s0 + s0, 4), dtype="float32") = R.concat((x,
y), axis=0)
+ gv: R.Tuple(R.Tensor((1 + s0 + s0, 4), dtype="float32")) =
(lv,)
R.output(gv)
return gv
@@ -8170,19 +8190,18 @@ def
test_dynamic_shape_with_multiplication_constraints():
class Expected:
@R.function
def main(
- x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("s0_2", 4),
dtype="float32")
- ) -> R.Tuple(R.Tensor(("s0 + s0_2", 4), dtype="float32")):
+ x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("2 * s0",
4), dtype="float32")
+ ) -> R.Tuple(R.Tensor(("s0 + 2 * s0", 4), dtype="float32")):
s0 = T.int64()
- s0_2 = T.int64()
R.func_attr(
{
- "tir_var_lower_bound": {"s77": 1, "s77_2": 2},
- "tir_var_upper_bound": {"s77": 64, "s77_2": 128},
+ "tir_var_lower_bound": {"s77": 1},
+ "tir_var_upper_bound": {"s77": 64},
}
)
with R.dataflow():
- lv: R.Tensor((s0 + s0_2, 4), dtype="float32") = R.concat((x,
y), axis=0)
- gv: R.Tuple(R.Tensor((s0 + s0_2, 4), dtype="float32")) = (lv,)
+ lv: R.Tensor((s0 + 2 * s0, 4), dtype="float32") = R.concat((x,
y), axis=0)
+ gv: R.Tuple(R.Tensor((s0 + 2 * s0, 4), dtype="float32")) =
(lv,)
R.output(gv)
return gv
