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 eeb1d61c83 [Fix][Relax][Torch] Materialize runtime scalar shape values
(#20138)
eeb1d61c83 is described below
commit eeb1d61c83a970fab10af32cdc2c800613c022ff
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Aug 29 21:20:26 2026 -0700
[Fix][Relax][Torch] Materialize runtime scalar shape values (#20138)
Materialize runtime integer scalars from `torch.export` as Relax
symbolic values and convert them back to rank-zero tensors for fill
operations. Preserve `full`/`full_like` dtype semantics, support direct
`fill`/`fill_` without decomposition, and lower symbolic shape
comparisons as real predicates. Runtime scalar extraction now also
accepts statically one-element integer tensors at higher ranks and
signed/unsigned integer dtypes, preserving negative, unsigned, and
64-bit values. Python scalar fills are materialized in the destination
dtype across `full`, `full_like`, `fill`, `fill_`, `masked_fill`, and
`masked_fill_`, avoiding 32-bit narrowing or precision loss. Symbolic
boolean fills, negation, floor division, and modulo are also supported.
---
.../frontend/torch/base_fx_graph_translator.py | 113 ++++--
.../frontend/torch/exported_program_translator.py | 25 +-
tests/python/relax/test_frontend_dynamo.py | 2 +-
.../relax/test_frontend_from_exported_program.py | 400 ++++++++++++++++++++-
4 files changed, 499 insertions(+), 41 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 b0bb14ac95..d600987cdd 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -2501,15 +2501,35 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _fill(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.ty.dtype
- value = args[1] if isinstance(args[1], relax.Expr) else
relax.const(args[1], dtype)
+ dtype = str(x.ty.dtype)
+ value = self._convert_scalar_fill_value(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 or boolean primitive value as a rank-zero
tensor."""
+ if not value.ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT,
DataTypeCode.BOOL):
+ raise TypeError(f"Cannot materialize primitive value of dtype
{value.ty} as a tensor")
+ shape_value = value if str(value.ty) == "int64" else
value.astype("int64")
+ value_tensor = self.block_builder.emit(
+ relax.op.shape_to_tensor(relax.ShapeExpr([shape_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 _convert_scalar_fill_value(self, value, dtype: str) -> relax.Expr:
+ """Convert a PyTorch scalar fill value to a rank-zero Relax tensor."""
+ if isinstance(getattr(value, "ty", None), PrimType):
+ return self._prim_value_to_scalar_tensor(value, dtype)
+ if isinstance(value, relax.Expr):
+ return value
+ return relax.const(value, dtype)
+
def _inplace_fill(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
x = args[0]
- dtype = x.ty.dtype.dtype
- value = args[1] if isinstance(args[1], relax.Expr) else
relax.const(args[1], dtype)
+ dtype = str(x.ty.dtype)
+ value = self._convert_scalar_fill_value(args[1], dtype)
filled = self.block_builder.emit(relax.op.full(x.ty.shape, value,
dtype))
self.env[node.args[0]] = filled
return filled
@@ -2519,10 +2539,23 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
args = self.retrieve_args(node)
size = relax.ShapeExpr(args[0] if isinstance(args[0], list | tuple)
else (args[0],))
- 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)
+ torch_dtype = node.kwargs.get("dtype")
+ if torch_dtype is None:
+ output_meta = node.meta.get("val")
+ if output_meta is None:
+ output_meta = node.meta.get("tensor_meta")
+ torch_dtype = getattr(output_meta, "dtype", None)
+ if torch_dtype is None:
+ if isinstance(args[1], bool):
+ torch_dtype = "bool"
+ elif isinstance(args[1], int):
+ torch_dtype = "int64"
+ elif isinstance(getattr(args[1], "ty", None), PrimType):
+ torch_dtype = args[1].ty.dtype
+ else:
+ torch_dtype = torch.get_default_dtype()
+ dtype = self._convert_data_type(torch_dtype, self.env)
+ value = self._convert_scalar_fill_value(args[1], dtype)
return self.block_builder.emit(
relax.op.full(
size,
@@ -2532,15 +2565,28 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
)
def _full_like(self, node: fx.Node) -> relax.Var:
- x = self.env[node.args[0]]
- value = node.args[1]
- 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)):
+ args = self.retrieve_args(node)
+ x = args[0]
+ value = args[1]
+ x_dtype = str(x.ty.dtype)
+ torch_dtype = node.kwargs.get("dtype")
+ dtype = self._convert_data_type(x_dtype if torch_dtype is None else
torch_dtype, self.env)
+ fill_dtype = dtype if dtype != x_dtype else None
+ if (
+ fill_dtype is None
+ and isinstance(value, int | float)
+ and (math.isinf(value) or math.isnan(value))
+ ):
if not ("float" in x_dtype or "bfloat16" in x_dtype):
fill_dtype = "float32"
+ dtype = fill_dtype
+
+ if isinstance(getattr(value, "ty", None), PrimType):
+ fill_value = self._prim_value_to_scalar_tensor(value, dtype)
+ elif isinstance(value, relax.Expr):
+ fill_value = value
+ else:
+ fill_value = relax.const(value, dtype)
return self.block_builder.emit(relax.op.full_like(x, fill_value,
dtype=fill_dtype))
@@ -2551,16 +2597,14 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return self.block_builder.emit(relax.op.take(x, index, dim))
def _inplace_masked_fill(self, node: fx.Node) -> relax.Var:
- x = self.env[node.args[0]]
- mask = self.env[node.args[1]]
- value = node.args[2]
- rx_value = relax.const(value)
-
- x_dtype = x.ty.dtype.dtype
+ args = self.retrieve_args(node)
+ x, mask, value = args[:3]
+ x_dtype = str(x.ty.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):
fill_dtype = "float32"
+ rx_value = self._convert_scalar_fill_value(value, fill_dtype or
x_dtype)
values = self.block_builder.emit(relax.op.full_like(x, rx_value,
dtype=fill_dtype))
@@ -2596,16 +2640,14 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
)
def _masked_fill(self, node: fx.Node) -> relax.Var:
- x = self.env[node.args[0]]
- mask = self.env[node.args[1]]
- value = node.args[2]
- rx_value = relax.const(value)
-
- x_dtype = x.ty.dtype.dtype
+ args = self.retrieve_args(node)
+ x, mask, value = args[:3]
+ x_dtype = str(x.ty.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):
fill_dtype = "float32"
+ rx_value = self._convert_scalar_fill_value(value, fill_dtype or
x_dtype)
values = self.block_builder.emit(relax.op.full_like(x, rx_value,
dtype=fill_dtype))
@@ -2807,6 +2849,23 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
def _item(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
+ shape = self.shape_of(x)
+ dtype = x.ty.dtype
+ analyzer = tvm.arith.Analyzer()
+ has_single_element = shape is not None and all(
+ analyzer.can_prove_equal(dim, 1) for dim in shape
+ )
+ if has_single_element and dtype.matches_code(DataTypeCode.INT,
DataTypeCode.UINT):
+ scalar = x
+ if str(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
+ if shape is not None and len(shape) == 0:
+ 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 0e69074af6..ced0aa7b28 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1481,8 +1481,12 @@ class ExportedProgramImporter(BaseFXGraphImporter):
########## Symbolic Shape Constraints ##########
- def _symbolic_comparison(self, _: fx.Node) -> relax.Expr:
- return self.block_builder.emit(relax.const(True, dtype="bool"))
+ def _symbolic_comparison(self, intrinsic_op: Callable) -> Callable:
+ def convert(node: fx.Node) -> relax.Expr:
+ lhs, rhs = self.retrieve_args(node)
+ return self.block_builder.emit(relax.prim_value(intrinsic_op(lhs,
rhs)))
+
+ return convert
########## Higher-Order Ops ##########
@@ -1754,6 +1758,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"logical_xor.default": self._logical_xor,
"log_softmax.int": self._log_softmax,
"_log_softmax.default": self._log_softmax,
+ "neg": lambda node: operator.neg(self.retrieve_args(node)[0]),
"neg.default": self._unary_op(relax.op.negative),
"pad.default": self._pad,
"constant_pad_nd.default": self._constant_pad_nd,
@@ -1792,6 +1797,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"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),
@@ -1810,6 +1816,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"eq.Scalar": self._binary_op(relax.op.equal, operator.eq),
"eq.Tensor": self._binary_op(relax.op.equal, operator.eq),
"floor_divide.default": self._binary_op(relax.op.floor_divide,
operator.floordiv),
+ "floordiv": self._binary_op(relax.op.floor_divide,
operator.floordiv),
"fmod.Scalar": self._fmod,
"fmod.Tensor": self._fmod,
"logaddexp.default": self._binary_op(relax.op.log_add_exp,
torch.logaddexp),
@@ -1835,6 +1842,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"minimum.default": self._binary_op(relax.op.minimum,
torch.minimum),
"remainder.Tensor": self._binary_op(relax.op.floor_mod,
operator.mod),
"remainder.Scalar": self._binary_op(relax.op.floor_mod,
operator.mod),
+ "mod": self._binary_op(relax.op.floor_mod, operator.mod),
"mul": self._binary_op(relax.op.multiply, operator.mul),
"mul.Tensor": self._binary_op(relax.op.multiply, operator.mul),
"mul.Scalar": self._binary_op(relax.op.multiply, operator.mul),
@@ -1847,6 +1855,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"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_),
@@ -2031,13 +2040,15 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"item.default": self._item,
"sym_size.int": self._sym_size_int,
"_local_scalar_dense.default": self._item,
- # symbolic shape constraints (no-ops for compilation)
+ # symbolic shape operations and constraints
"sym_constrain_range_for_size.default": lambda node:
self.env[node.args[0]],
"_assert_scalar.default": lambda node: self.env[node.args[0]],
- "ge": self._symbolic_comparison,
- "le": self._symbolic_comparison,
- "gt": self._symbolic_comparison,
- "lt": self._symbolic_comparison,
+ "ge": self._symbolic_comparison(operator.ge),
+ "le": self._symbolic_comparison(operator.le),
+ "gt": self._symbolic_comparison(operator.gt),
+ "lt": self._symbolic_comparison(operator.lt),
+ "eq": self._symbolic_comparison(operator.eq),
+ "ne": self._symbolic_comparison(operator.ne),
# higher-order ops
"cond": self._cond,
}
diff --git a/tests/python/relax/test_frontend_dynamo.py
b/tests/python/relax/test_frontend_dynamo.py
index 72e05a5649..e795b0881f 100644
--- a/tests/python/relax/test_frontend_dynamo.py
+++ b/tests/python/relax/test_frontend_dynamo.py
@@ -504,7 +504,7 @@ def test_masked_fill():
) -> R.Tensor((256, 256), dtype="float32"):
with R.dataflow():
lv: R.Tensor((256, 256), dtype="float32") = R.full_like(
- inp_1, R.const(0, "int32"), dtype=None
+ inp_1, R.const(0.0, "float32"), dtype=None
)
lv1: R.Tensor((256, 256), dtype="float32") = R.where(inp_0,
lv, inp_1)
gv: R.Tensor((256, 256), dtype="float32") = lv1
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index e9d2ac8b70..7dc3c73564 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -5549,6 +5549,265 @@ def test_expand_with_new_leading_dimension():
assert tvm.arith.Analyzer().can_prove_equal(output_shape[2],
input_shape[1])
+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())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
[email protected](
+ ("item_dtype", "item_shape"),
+ [
+ (torch.int8, ()),
+ (torch.uint8, ()),
+ (torch.int16, ()),
+ (torch.int32, ()),
+ (torch.int64, ()),
+ (torch.int64, (1,)),
+ (torch.int64, (1, 1)),
+ ],
+)
+def test_dynamic_scalar_item_single_element_integer_dtypes(item_dtype,
item_shape):
+ class DynamicItem(torch.nn.Module):
+ def forward(self, x):
+ lengths = torch.full(
+ (x.shape[0],),
+ x.shape[1],
+ device=x.device,
+ dtype=item_dtype,
+ )
+ value = lengths.max().reshape(item_shape).item()
+ return torch.arange(value, device=x.device), torch.full_like(
+ x, value, dtype=torch.int64
+ )
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ example_args = (torch.randn(3, 4, dtype=torch.float32),)
+ exported_program = export(
+ DynamicItem(),
+ args=example_args,
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape in ((3, 4), (5, 2)):
+ torch_input = torch.randn(shape, dtype=torch.float32)
+ expected = DynamicItem()(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())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_runtime_scalar_item_integer_value_ranges():
+ class RuntimeItems(torch.nn.Module):
+ def forward(self, x, signed, unsigned, wide):
+ return (
+ torch.full_like(x, signed.item(), dtype=torch.int64),
+ torch.full_like(x, unsigned.item(), dtype=torch.int64),
+ torch.full_like(x, wide.item(), dtype=torch.int64),
+ )
+
+ example_args = (
+ torch.randn(2, 3, dtype=torch.float32),
+ torch.tensor([[-3]], dtype=torch.int8),
+ torch.tensor([200], dtype=torch.uint8),
+ torch.tensor(1 << 40, dtype=torch.int64),
+ )
+ exported_program = export(RuntimeItems(), args=example_args)
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ expected = RuntimeItems()(*example_args)
+ actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in
example_args))
+ for actual_value, expected_value in zip(actual, expected):
+ np.testing.assert_array_equal(actual_value.numpy(),
expected_value.numpy())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_dynamic_scalar_fill_operations():
+ class DynamicFills(torch.nn.Module):
+ def forward(self, x):
+ inferred_dtype = torch.full((x.shape[0],), x.shape[1],
device=x.device)
+ explicit_dtype = torch.full_like(x, x.shape[0],
dtype=torch.float64)
+ filled = torch.fill(x, x.shape[1])
+ filled_inplace = x.clone()
+ filled_inplace.fill_(x.shape[0])
+ return inferred_dtype, explicit_dtype, filled, filled_inplace
+
+ example_args = (torch.randn(3, 4, dtype=torch.float32),)
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ exported_program = export(
+ DynamicFills(),
+ args=example_args,
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape in ((3, 4), (5, 2)):
+ torch_input = torch.randn(shape, dtype=torch.float32)
+ expected = DynamicFills()(torch_input)
+ actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))
+ actual_arrays = [value.numpy() for value in actual]
+
+ assert actual_arrays[0].dtype == np.dtype("int64")
+ assert actual_arrays[1].dtype == np.dtype("float64")
+ for actual_value, expected_value in zip(actual_arrays, expected):
+ assert actual_value.dtype == expected_value.numpy().dtype
+ np.testing.assert_array_equal(actual_value, expected_value.numpy())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_dynamic_boolean_fill_operations():
+ class DynamicBooleanFills(torch.nn.Module):
+ def forward(self, x):
+ value = x.shape[0] == x.shape[1]
+ base = x > 0
+ filled = torch.fill(base, value)
+ filled_inplace = base.clone()
+ filled_inplace.fill_(value)
+ created = torch.full(
+ (x.shape[0], x.shape[1]),
+ value,
+ device=x.device,
+ dtype=torch.bool,
+ )
+ like = torch.full_like(x, value, dtype=torch.bool)
+ return created, like, filled, filled_inplace
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ example_args = (torch.randn(3, 4, dtype=torch.float32),)
+ exported_program = export(
+ DynamicBooleanFills(),
+ args=example_args,
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape in ((3, 3), (3, 4)):
+ torch_input = torch.randn(shape, dtype=torch.float32)
+ expected = DynamicBooleanFills()(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())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_dynamic_scalar_arithmetic():
+ class DynamicScalarArithmetic(torch.nn.Module):
+ def forward(self, x):
+ lengths = torch.full(
+ (x.shape[0],),
+ x.shape[1],
+ device=x.device,
+ dtype=torch.int64,
+ )
+ value = lengths.max().item()
+ return (
+ torch.full_like(x, -value, dtype=torch.int64),
+ torch.arange(value // 2, device=x.device),
+ torch.arange(value % 3, device=x.device),
+ )
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ example_args = (torch.randn(3, 4, dtype=torch.float32),)
+ exported_program = export(
+ DynamicScalarArithmetic(),
+ args=example_args,
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape in ((3, 4), (5, 2)):
+ torch_input = torch.randn(shape, dtype=torch.float32)
+ expected = DynamicScalarArithmetic()(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())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
[email protected](
+ ("fill_value", "dtype"),
+ [
+ (1 << 40, torch.int64),
+ (1.0 + 2**-40, torch.float64),
+ ],
+)
+def test_full_and_full_like_python_fill_value_uses_explicit_dtype(fill_value,
dtype):
+ class FullLike(torch.nn.Module):
+ def forward(self, x):
+ return (
+ torch.full(x.shape, fill_value, dtype=dtype),
+ torch.full_like(x, fill_value, dtype=dtype),
+ )
+
+ example_args = (torch.randn(2, 3, dtype=torch.float32),)
+ exported_program = export(FullLike(), args=example_args)
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ expected = FullLike()(*example_args)
+ actual = vm["main"](tvm.runtime.tensor(example_args[0].numpy()))
+ for actual_value, expected_value in zip(actual, expected):
+ assert actual_value.numpy().dtype == expected_value.numpy().dtype
+ np.testing.assert_array_equal(actual_value.numpy(),
expected_value.numpy())
+
+
def test_split():
class Chunk(Module):
def forward(self, input):
@@ -6072,6 +6331,71 @@ def test_masked_fill_inplace():
verify_model(Masked_Fill_Inplace(), example_args, {}, Expected)
[email protected](not env.has_llvm(), reason="need llvm")
+def test_dynamic_scalar_masked_fill_operations():
+ class DynamicMaskedFills(torch.nn.Module):
+ def forward(self, x):
+ mask = x > 0
+ filled = x.masked_fill(mask, x.shape[0])
+ filled_inplace = x.clone()
+ filled_inplace.masked_fill_(mask, x.shape[1])
+ return filled, filled_inplace
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ example_args = (torch.randn(3, 4, dtype=torch.float32),)
+ exported_program = export(
+ DynamicMaskedFills(),
+ args=example_args,
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape in ((3, 4), (5, 2)):
+ torch_input = torch.randn(shape, dtype=torch.float32)
+ expected = DynamicMaskedFills()(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())
+
+
[email protected](not env.has_llvm(), reason="need llvm")
[email protected](
+ ("fill_value", "dtype"),
+ [
+ (1 << 40, torch.int64),
+ (1.0 + 2**-40, torch.float64),
+ ],
+)
+def test_fill_and_masked_fill_python_value_uses_input_dtype(fill_value, dtype):
+ class MaskedFills(torch.nn.Module):
+ def forward(self, x, mask):
+ filled_all = torch.fill(x, fill_value)
+ filled_all_inplace = x.clone()
+ filled_all_inplace.fill_(fill_value)
+ filled = x.masked_fill(mask, fill_value)
+ filled_inplace = x.clone()
+ filled_inplace.masked_fill_(mask, fill_value)
+ return filled_all, filled_all_inplace, filled, filled_inplace
+
+ example_args = (
+ torch.arange(6, dtype=dtype).reshape(2, 3),
+ torch.tensor([[True, False, True], [False, True, False]]),
+ )
+ exported_program = export(MaskedFills(), args=example_args)
+ mod = from_exported_program(exported_program, run_ep_decomposition=False)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ expected = MaskedFills()(*example_args)
+ actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in
example_args))
+ for actual_value, expected_value in zip(actual, expected):
+ assert actual_value.numpy().dtype == expected_value.numpy().dtype
+ np.testing.assert_array_equal(actual_value.numpy(),
expected_value.numpy())
+
+
def test_masked_select():
class MaskedSelect(Module):
def forward(self, data: torch.Tensor, mask: torch.Tensor):
@@ -6094,8 +6418,8 @@ def test_masked_select():
)
lv4: R.Tensor((u0,), dtype="int64") = R.squeeze(lv3, axis=[0])
lv5: R.Tensor((u0,), dtype="float32") = R.take(lv, lv4,
axis=0, mode="fast")
- lv6: R.Tensor((), dtype="bool") = R.const(True, "bool")
- lv7: R.Tensor((), dtype="bool") = R.const(True, "bool")
+ lv6: T.bool = u0 >= 0
+ lv7: T.bool = u0 <= 6
gv: R.Tuple(R.Tensor((u0,), dtype="float32")) = (lv5,)
R.output(gv)
return gv
@@ -6430,7 +6754,9 @@ def test_ones_like():
R.Tensor((128, 128), dtype="float32")
):
with R.dataflow():
- lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input,
R.const(1, "int32"))
+ lv: R.Tensor((128, 128), dtype="float32") = R.full_like(
+ input, R.const(1.0, "float32")
+ )
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
R.output(gv)
return gv
@@ -6452,7 +6778,9 @@ def test_zero_inplace():
R.Tensor((128, 128), dtype="float32")
):
with R.dataflow():
- lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input,
R.const(0, "int32"))
+ lv: R.Tensor((128, 128), dtype="float32") = R.full_like(
+ input, R.const(0.0, "float32")
+ )
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
R.output(gv)
return gv
@@ -6498,7 +6826,9 @@ def test_zeros_like():
R.Tensor((128, 128), dtype="float32")
):
with R.dataflow():
- lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input,
R.const(0, "int32"))
+ lv: R.Tensor((128, 128), dtype="float32") = R.full_like(
+ input, R.const(0.0, "float32")
+ )
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
R.output(gv)
return gv
@@ -8692,7 +9022,7 @@ def test_cond_shape_predicate():
s77 = T.int64()
R.func_attr({"tir_var_lower_bound": {"s77": 1}})
cls = expected
- gv: R.Tensor((), dtype="bool") = R.const(True, "bool")
+ gv: T.bool = s77 > 4
if gv:
gv1: R.Tensor((s77, 4), dtype="float32") =
cls.cond_true_branch_0(x)
cond_result: R.Tensor((s77, 4), dtype="float32") = gv1
@@ -8712,6 +9042,64 @@ def test_cond_shape_predicate():
)
[email protected](not env.has_llvm(), reason="need llvm")
+def test_cond_shape_equality_predicate():
+ class CondShapeEqualityModel(Module):
+ def forward(self, x):
+ def true_fn(x):
+ return x + 1.0
+
+ def false_fn(x):
+ return x - 1.0
+
+ return torch.cond(x.shape[0] == x.shape[1], true_fn, false_fn,
(x,))
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ exported_program = export(
+ CondShapeEqualityModel(),
+ args=(torch.zeros(3, 3),),
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape, expected_value in (((3, 3), 1.0), ((2, 3), -1.0)):
+ torch_input = torch.zeros(shape, dtype=torch.float32)
+ actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0]
+ np.testing.assert_array_equal(actual.numpy(), np.full(shape,
expected_value, "float32"))
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_cond_shape_inequality_predicate():
+ class CondShapeInequalityModel(Module):
+ def forward(self, x):
+ def true_fn(x):
+ return x + 1.0
+
+ def false_fn(x):
+ return x - 1.0
+
+ return torch.cond(x.shape[0] != x.shape[1], true_fn, false_fn,
(x,))
+
+ rows = torch.export.Dim("rows", min=1, max=8)
+ columns = torch.export.Dim("columns", min=1, max=8)
+ exported_program = export(
+ CondShapeInequalityModel(),
+ args=(torch.zeros(2, 3),),
+ dynamic_shapes={"x": {0: rows, 1: columns}},
+ )
+ mod = from_exported_program(exported_program)
+ executable = relax.build(mod, tvm.target.Target("llvm"))
+ vm = relax.VirtualMachine(executable, tvm.cpu())
+
+ for shape, expected_value in (((2, 3), 1.0), ((3, 3), -1.0)):
+ torch_input = torch.zeros(shape, dtype=torch.float32)
+ actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0]
+ np.testing.assert_array_equal(actual.numpy(), np.full(shape,
expected_value, "float32"))
+
+
def test_cond_tuple_output():
"""Cond where both branches return a tuple."""
