This is an automated email from the ASF dual-hosted git repository.
tqchen 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 0c2b1029d5 [Frontend][PyTorch] Simplify tests and support exported
assertions (#20360)
0c2b1029d5 is described below
commit 0c2b1029d5e785cd3348d6b9c9ddac9b5a7e462b
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Sep 16 07:44:44 2026 -0400
[Frontend][PyTorch] Simplify tests and support exported assertions (#20360)
This pr simplifies duplicate tests across the FX and ExportedProgram
frontends, reuse equivalent models and expected IR, and trim redundant
parameter cases. This reduces test code by 806 lines while retaining
representative operator coverage.
Support `aten._assert_async` by lowering it to Relax runtime assertions,
including assertions inside `torch.cond` branches. This fixes `one_hot`
import failures with PyTorch 2.12 while preserving invalid-index checks
and error messages.
Also apply Ruff formatting to three existing CUDA files to resolve the
lint failure.
---
.../frontend/torch/exported_program_translator.py | 41 +-
.../relax/test_frontend_from_exported_program.py | 780 +++------------------
tests/python/relax/test_frontend_from_fx.py | 308 +-------
3 files changed, 177 insertions(+), 952 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index f3bfce230a..dbf4993251 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1647,7 +1647,9 @@ class ExportedProgramImporter(BaseFXGraphImporter):
self.env[ph] = param
# Build the branch function (using a plain BindingBlock, not
DataflowBlock).
- with self.block_builder.function(name=unique_name, params=params):
+ with self.block_builder.function(
+ name=unique_name, params=params, pure=not
self._has_assert_async(graph_module)
+ ):
inner = self._translate_fx_graph(graph_module, nodes, {})
if isinstance(inner, tuple | list):
if len(inner) == 1:
@@ -1710,6 +1712,32 @@ class ExportedProgramImporter(BaseFXGraphImporter):
########## Others ##########
+ @staticmethod
+ def _has_assert_async(graph_module) -> bool:
+ return any(
+ node.op == "call_function"
+ and node.target
+ in (torch.ops.aten._assert_async.default,
torch.ops.aten._assert_async.msg)
+ for module in graph_module.modules()
+ if isinstance(module, fx.GraphModule)
+ for node in module.graph.nodes
+ )
+
+ def _assert_async(self, node: fx.Node) -> relax.Var:
+ condition = self.env[node.args[0]]
+ if condition.ty.dtype.dtype != "bool":
+ condition = self.block_builder.emit(relax.op.astype(condition,
"bool"))
+ if condition.ty.ndim != 0:
+ condition = self.block_builder.emit(relax.op.reshape(condition,
[]))
+ message = (
+ node.args[1]
+ if len(node.args) > 1
+ else node.kwargs.get("assert_msg", "Assertion Failed")
+ )
+ return self.block_builder.emit(
+ relax.op.assert_op(condition, [relax.StringImm(message)],
format="{}")
+ )
+
def create_convert_map(
self,
) -> dict[str, Callable[[fx.Node], relax.Var]]:
@@ -2049,6 +2077,8 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"to.dtype_layout": self._to,
"type_as.default": self._type_as,
# other
+ "_assert_async.default": self._assert_async,
+ "_assert_async.msg": self._assert_async,
"getitem": self._getitem,
"item.default": self._item,
"sym_size.int": self._sym_size_int,
@@ -2242,12 +2272,13 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# Find all the missing function types
self._check_unsupported_func_type(nodes)
- # When the graph contains torch.cond, we must avoid DataflowBlock
- # because relax.If cannot appear inside a dataflow region.
- use_dataflow = not self._has_cond_op(nodes)
+ # Assertions have side effects, including when they occur in a cond
branch.
+ # Neither these effects nor relax.If may appear inside a dataflow
region.
+ is_pure = not self._has_assert_async(exported_program.graph_module)
+ use_dataflow = is_pure and not self._has_cond_op(nodes)
with self.block_builder.function(
- name=func_name, params=list(inputs_vars.values()).copy(),
attrs=func_attrs
+ name=func_name, params=list(inputs_vars.values()).copy(),
attrs=func_attrs, pure=is_pure
):
with contextlib.ExitStack() as stack:
if use_dataflow:
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 484c68bc66..d78629f573 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -148,18 +148,8 @@ def test_basic_unary_ops(pytorch_op, relax_op):
verify_model(UnaryOp(), example_args, {}, expected)
-def test_round_decimals():
- """torch.round(x, decimals) is exported as aten.round.decimals, which was
missing
- from the convert map (only round.default was registered) and made any
explicit
- decimals -- including decimals=0 -- fail with
- "AssertionError: Unsupported function types ['round.decimals']".
-
- With the decimals overload registered, torch.round(x, decimals) must
convert and
- match PyTorch's round-half-to-even results, including negative decimals
- (round(25, -1) == 20) where the scale-by-0.1 float precision path used to
be wrong.
- """
-
- class RoundDecimalsModel(Module):
+def test_round():
+ class Round(Module):
def __init__(self, decimals):
super().__init__()
self.decimals = decimals
@@ -167,77 +157,13 @@ def test_round_decimals():
def forward(self, input):
return torch.round(input, decimals=self.decimals)
- # Half values exercise ties-to-even; 25/125/165 exercise the
negative-decimals path.
+ # Ties-to-even, negative decimals, and scales that overflow to NaN.
x = torch.tensor(
- [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25],
dtype=torch.float32
+ [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25, 0.0],
+ dtype=torch.float32,
)
- for decimals in (0, 1, -1, -2):
- verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,),
rtol=1e-6, atol=1e-6)
-
-
-def test_round_decimals_low_precision():
- """Scaling for low-precision inputs must happen in float32 and be cast
back.
-
- 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 ==
250000
- exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5
- already overflows float16 (the scale itself becomes inf), turning
decimals=5
- and -5 into NaN. Upcasting the input to float32 keeps the scaling exact;
the
- rounded result is cast back to the input dtype.
- """
-
- class RoundDecimalsModel(Module):
- def __init__(self, decimals):
- super().__init__()
- self.decimals = decimals
-
- def forward(self, input):
- return torch.round(input, decimals=self.decimals)
-
- x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5],
dtype=torch.float16)
- # Positive decimals exercise the multiply-by-10**d overflow (4, 5);
- # negative decimals exercise the 10**|d| scale overflowing float16 (-5).
- for decimals in (2, 4, 5, -2, -4, -5):
- verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,),
rtol=1e-6, atol=1e-6)
-
-
-def test_round_decimals_large():
- """A large |decimals| must import and run without OverflowError.
-
- The scale 10**|decimals| used to be built as an unbounded host Python int
- before being handed to relax.const, whose int-to-float conversion raises
- OverflowError ("int too large to convert to float") once |decimals| >= 309
- (10**309 already exceeds the float64 range). PyTorch accepts such decimals
and
- exports a valid aten.round.decimals node, so importing the exported program
- must not crash on them. The scale is now built directly in the float dtype
and
- saturates to inf once it leaves the finite range, matching PyTorch, whose
- all-NaN result here comes from the same inf scale.
- """
-
- class RoundDecimalsModel(Module):
- def __init__(self, decimals):
- super().__init__()
- self.decimals = decimals
-
- def forward(self, input):
- return torch.round(input, decimals=self.decimals)
-
- x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32)
- for decimals in (309, -309):
- exported_program = export(RoundDecimalsModel(decimals).eval(),
args=(x,))
- mod = from_exported_program(exported_program) # used to raise
OverflowError here
- ex = relax.build(mod, target="llvm")
- vm = relax.VirtualMachine(ex, tvm.cpu())
- tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
- got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
-
- # The scale overflows to inf, and IEEE arithmetic turns every element
into
- # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
- # (empty here) finite elements separately, since allclose fails on NaN.
- expected = torch.round(x, decimals=decimals)
- actual = torch.as_tensor(got)
- assert torch.equal(torch.isnan(actual), torch.isnan(expected))
- finite = ~torch.isnan(expected)
- assert torch.allclose(actual[finite], expected[finite], rtol=1e-6,
atol=1e-6)
+ for decimals in (0, 1, -1, -2, 309, -309):
+ verify_model_numerically(Round(decimals).eval(), (x,), rtol=1e-6,
atol=1e-6)
operator_bool_unary = [
@@ -523,7 +449,7 @@ def test_extended_unary_ops():
return torch.ops.aten.hardswish_(input)
@tvm.script.ir_module
- class expected_hardswish_for_1_2:
+ class expected_hardswish:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((1, 3, 10, 10), dtype="float32")
@@ -546,32 +472,8 @@ def test_extended_unary_ops():
R.output(gv)
return gv
- @tvm.script.ir_module
- class expected_hardswish_for_3:
- @R.function
- def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 3, 10, 10), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
- input, R.const(3.0, "float32")
- )
- lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
- lv, R.prim_value(0), R.prim_value(T.float64("inf"))
- )
- lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
- lv1, R.prim_value(T.float64("-inf")), R.prim_value(6)
- )
- lv3: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.multiply(input, lv2)
- lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
- lv3, R.const(6.0, "float32")
- )
- gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv4,)
- R.output(gv)
- return gv
-
- verify_model(Hardswish(), example_args, {}, expected_hardswish_for_1_2)
- verify_model(Hardswish3(), example_args, {}, expected_hardswish_for_3)
+ verify_model(Hardswish(), example_args, {}, expected_hardswish)
+ verify_model(Hardswish3(), example_args, {}, expected_hardswish)
# isfinite
class IsFinite(Module):
@@ -835,20 +737,7 @@ def test_extended_unary_ops():
def forward(self, input):
return torch.ops.aten.silu_(input)
- @tvm.script.ir_module
- class expected_silu_:
- @R.function
- def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 3, 10, 10), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.sigmoid(input)
- lv1: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.multiply(input, lv)
- gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
- R.output(gv)
- return gv
-
- verify_model(SiLU_(), example_args, {}, expected_silu_)
+ verify_model(SiLU_(), example_args, {}, expected_silu)
# square
class Square(Module):
@@ -2033,16 +1922,10 @@ def test_adaptive_avgpool3d():
def test_addmm():
class Addmm1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3)
class Addmm2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
@@ -2092,16 +1975,10 @@ def test_addmm():
def test_sparse_addmm():
class SparseAddmm1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.sparse.addmm(x1, x2, x3)
class SparseAddmm2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.sparse.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
@@ -2447,9 +2324,6 @@ def test_avg_pool3d():
def test_baddbmm():
class BAddBMM1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, c, x, y):
return torch.baddbmm(c, x, y)
@@ -2471,9 +2345,6 @@ def test_baddbmm():
return gv
class BAddBMM2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, c, x, y):
return torch.baddbmm(c, x, y, alpha=2, beta=0)
@@ -2497,9 +2368,6 @@ def test_baddbmm():
return gv
class BAddBMM3(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, c, x, y):
return torch.baddbmm(c, x, y, alpha=2, beta=3)
@@ -2558,9 +2426,6 @@ def test_baddbmm():
def test_bmm():
class BMM(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return torch.bmm(x, y)
@@ -3372,19 +3237,13 @@ def test_pixel_shuffle():
def test_einsum():
- class Einsum1(Module):
- def __init__(self):
+ class Einsum(Module):
+ def __init__(self, subscripts):
super().__init__()
+ self.subscripts = subscripts
- def forward(self, x):
- return torch.einsum("ii", x)
-
- class Einsum2(Module):
- def __init__(self):
- super().__init__()
-
- def forward(self, x, y):
- return torch.einsum("i,j->ij", x, y)
+ def forward(self, *args):
+ return torch.einsum(self.subscripts, *args)
@tvm.script.ir_module
class Expected1:
@@ -3413,34 +3272,12 @@ def test_einsum():
return gv
example_args = (torch.randn(4, 4, dtype=torch.float32),)
- verify_model(Einsum1(), example_args, {}, Expected1,
run_ep_decomposition=False)
+ verify_model(Einsum("ii"), example_args, {}, Expected1,
run_ep_decomposition=False)
example_args = (torch.randn(5, dtype=torch.float32), torch.randn(4,
dtype=torch.float32))
- verify_model(Einsum2(), example_args, {}, Expected2,
run_ep_decomposition=False)
-
-
-def test_einsum_repeated_subscript():
- """einsum with repeated subscripts (diagonal / trace) on the default
- decomposition path.
-
- ``run_decompositions`` (default) lowers repeated-subscript einsum to
- ``aten.diagonal`` + ``permute`` (+ ``sum`` for the trace), which the
- frontend converts with the ``_diagonal`` lowering. For the zero-offset
- square case (e.g. ``torch.einsum("ii->i")`` on an ``N x N`` input) the
- frontend emits a single repeated-subscript einsum that reads the diagonal
- directly; otherwise it permutes the diagonal dims to the trailing two axes,
- slices each to the diagonal length, and runs an einsum ``...zz->...z``.
- This used to raise ``AssertionError: Unsupported function types
- ['diagonal.default']``.
- """
-
- class EinsumDiag(Module):
- def __init__(self):
- super().__init__()
-
- def forward(self, x):
- return torch.einsum("ii->i", x)
+ verify_model(Einsum("i,j->ij"), example_args, {}, Expected2,
run_ep_decomposition=False)
+ # Default decomposition lowers repeated subscripts through aten.diagonal.
@tvm.script.ir_module
class Expected:
@R.function
@@ -3454,28 +3291,19 @@ def test_einsum_repeated_subscript():
return gv
example_args = (torch.randn(3, 3, dtype=torch.float32),)
- verify_model(EinsumDiag(), example_args, {}, Expected)
+ verify_model(Einsum("ii->i"), example_args, {}, Expected)
- class TraceEinsum(Module):
- def forward(self, x):
- return torch.einsum("ii->", x)
-
- class BatchedDiagEinsum(Module):
- def forward(self, x):
- return torch.einsum("...ii->...i", x)
-
- class AttentionEinsum(Module):
- def forward(self, x, y):
- return torch.einsum("abca,abcb->c", x, y)
-
- verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),))
- verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),))
- verify_model_numerically(AttentionEinsum(), (torch.randn(3, 3, 4, 3),
torch.randn(3, 3, 4, 3)))
+ for subscripts, shapes in [
+ ("ii->", [(4, 4)]),
+ ("...ii->...i", [(2, 3, 3)]),
+ ("abca,abcb->c", [(3, 3, 4, 3), (3, 3, 4, 3)]),
+ ]:
+ verify_model_numerically(Einsum(subscripts), tuple(torch.randn(*shape)
for shape in shapes))
class DirectDiagonal(Module):
- def __init__(self):
+ def __init__(self, offset):
super().__init__()
- self.offset = 1
+ self.offset = offset
def forward(self, x):
return torch.diagonal(x, self.offset, 0, 1)
@@ -3484,48 +3312,16 @@ def test_einsum_repeated_subscript():
def forward(self, x):
return torch.trace(x)
- verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),))
+ # For a 3x4 input, 4 and -3 are the first empty diagonals. Larger offsets
+ # in either direction check that negative diagonal lengths are clamped to
zero.
+ for offset in [1, 4, 6, -3, -6]:
+ verify_model_numerically(DirectDiagonal(offset), (torch.randn(3, 4),))
verify_model_numerically(DirectTrace(), (torch.randn(4, 4),))
- # Out-of-range offsets (|offset| >= max(extent1, extent2)) are valid in
- # PyTorch and yield an empty diagonal of shape (0,); the lowering must
- # clamp the diagonal length to zero instead of producing negative slice
- # extents or a wrong non-empty shape.
- class DirectDiagonalOutOfRange(Module):
- def __init__(self, offset):
- super().__init__()
- self.offset = offset
-
- def forward(self, x):
- return torch.diagonal(x, self.offset, 0, 1)
-
- for offset in [4, 5, 6, -3, -4, -5, -6]:
- verify_model_numerically(DirectDiagonalOutOfRange(offset),
(torch.randn(3, 4),))
-
-
-def test_einsum_diagonal_lowers_without_full_size_intermediate():
- """Regression test: a zero-offset square diagonal must not materialize
- full-size intermediates.
-
- ``torch.einsum("ii->i")`` on an ``N x N`` input is decomposed to
- ``aten.diagonal`` by ``run_decompositions``. Lowering that diagonal by
- permuting the diagonal dims to the trailing axes, slicing each to the
- diagonal length, and running the ``...zz->...z`` einsum materializes three
- full-size ``N x N`` intermediates (an identity permute and two identity
- strided slices) and hence three O(N^2) copy loops before the final O(N)
- diagonal loop. The ``_diagonal`` fast path instead emits a single
- repeated-subscript einsum that reads the diagonal directly, so no full-size
- intermediate exists in the frontend graph (and therefore neither in the
- lowered TIR). Assert that every intermediate produced by a call is at most
- O(N), both before and after legalization.
- """
-
- class EinsumDiag(Module):
- def forward(self, x):
- return torch.einsum("ii->i", x)
-
+ # A square diagonal must read the input directly without O(N^2) copies,
+ # both before and after legalization.
n = 8
- exported_program = export(EinsumDiag(), args=(torch.randn(n, n),))
+ exported_program = export(Einsum("ii->i"), args=(torch.randn(n, n),))
mod = from_exported_program(exported_program)
def rank2_call_results(ir_mod):
@@ -3835,9 +3631,6 @@ def test_linear():
def test_maxpool1d():
class MaxPool1d(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool1d(input, kernel_size=2)
@@ -3924,9 +3717,6 @@ def test_maxpool2d():
return self.pool(input)
class MaxPool2d_functional(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1])
@@ -4042,9 +3832,6 @@ def test_maxpool3d():
return self.pool(input)
class MaxPool3d_functional(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1])
@@ -5118,16 +4905,10 @@ def test_argmax_argmin():
example_args = (torch.randn(256, 256, dtype=torch.float32),)
class Argmax1(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmax(input, dim=-1)
class Argmax2(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmax(input, dim=-1, keepdim=True)
@@ -5159,16 +4940,10 @@ def test_argmax_argmin():
verify_model(Argmax2(), example_args, {}, expected_argmax2)
class Argmin1(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmin(input)
class Argmin2(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmin(input, keepdim=True)
@@ -5315,15 +5090,6 @@ def test_flatten():
verify_model(Flatten(), example_args, {}, expected1)
-def test_flatten_zero_sized_dim():
- class Flatten(Module):
- def forward(self, x):
- return torch.flatten(x)
-
- verify_model_numerically(Flatten(), (torch.randn(2, 0, 4,
dtype=torch.float32),))
- verify_model_numerically(Flatten(), (torch.randn(2, 3, 0,
dtype=torch.float32),))
-
-
def test_meshgrid():
class Meshgrid1(Module):
def forward(self, input1, input2):
@@ -5498,43 +5264,6 @@ def test_reshape_zero_sized_dim():
verify_model_numerically(ReshapeTrailing(), (torch.randn(0, 3,
dtype=torch.float32),))
-def test_reshape_multiple_zero_sized_dims():
- # A literal zero only survives relax's copy rule at a position whose input
dimension is
- # itself zero, so targets holding several zeros need more than one
position rewritten.
- class TwoZeros(Module):
- def forward(self, x):
- return x.reshape(0, 0)
-
- class ThreeZeros(Module):
- def forward(self, x):
- return x.reshape(0, 0, 0)
-
- class ZeroPastInputRank(Module):
- def forward(self, x):
- return x.reshape(0, 0, 4)
-
- verify_model_numerically(TwoZeros(), (torch.randn(0, 3,
dtype=torch.float32),))
- verify_model_numerically(TwoZeros(), (torch.randn(3, 0,
dtype=torch.float32),))
- verify_model_numerically(ThreeZeros(), (torch.randn(0, 3, 5,
dtype=torch.float32),))
- verify_model_numerically(ZeroPastInputRank(), (torch.randn(2, 0, 4,
dtype=torch.float32),))
-
-
-def test_reshape_zero_sized_dim_dynamic_batch():
- # One statically known zero fixes the element count at zero whatever the
symbolic
- # dimension turns out to be, so the literal zero in the target still has
to survive.
- # Reading it as "copy the batch" gives a non-empty shape that torch never
produces.
- class Reshape(Module):
- def forward(self, x):
- return x.reshape(0, 4)
-
- batch = torch.export.Dim("batch", min=1, max=64)
- verify_model_numerically(
- Reshape(),
- (torch.randn(3, 0, 4, dtype=torch.float32),),
- dynamic_shapes={"x": {0: batch}},
- )
-
-
def test_reshape_zero_sized_dim_symbolic_target():
# Only a literal is read as "copy the input dimension", so a symbolic
dimension in the
# target is not one: it is carried through, and the literal zero beside it
still has to
@@ -5919,53 +5648,6 @@ def test_dynamic_scalar_item_in_shape_operations():
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())
-
-
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_runtime_scalar_item_integer_value_ranges():
class RuntimeItems(torch.nn.Module):
@@ -6229,10 +5911,6 @@ def test_split_int_split_size():
for shape, split_size, dim in [
((10,), 6, 0),
- ((10,), 7, 0),
- ((10,), 8, 0),
- ((10,), 9, 0),
- ((12,), 7, 0),
((12, 8), 5, 1),
((3, 10), 6, -1),
]:
@@ -6620,27 +6298,6 @@ def test_empty():
verify_model(Empty(), example_args, {}, Expected)
-def test_empty_without_dtype():
- class EmptyWithoutDtype(Module):
- def forward(self, input):
- return torch.empty((5, 5))
-
- @tvm.script.ir_module
- class Expected:
- @R.function
- def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
- R.Tensor((5, 5), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((5, 5), dtype="float32") = R.zeros(R.shape([5,
5]), dtype="float32")
- gv: R.Tuple(R.Tensor((5, 5), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- example_args = (torch.randn(10, 10, dtype=torch.float32),)
- verify_model(EmptyWithoutDtype(), example_args, {}, Expected)
-
-
def test_fill():
class Fill(Module):
def forward(self, input: torch.Tensor):
@@ -7000,9 +6657,6 @@ def test_keep_params():
def test_unwrap_unit_return_tuple():
class Identity(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x):
return (x,)
@@ -7023,9 +6677,6 @@ def test_unwrap_unit_return_tuple():
def test_no_bind_return_tuple():
class Identity(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return (x, y)
@@ -7120,6 +6771,7 @@ def test_empty_like():
verify_model(EmptyLike(), example_args, {}, Expected)
[email protected](not env.has_llvm(), reason="need llvm")
def test_one_hot():
class OneHot(Module):
def forward(self, indices):
@@ -7132,19 +6784,17 @@ def test_one_hot():
indices: R.Tensor((5,), dtype="int64"),
) -> R.Tuple(R.Tensor((5, 10), dtype="int64")):
with R.dataflow():
- lv: R.Tensor((10,), dtype="int64") = R.arange(
- R.prim_value(0), R.prim_value(10), R.prim_value(1),
dtype="int64"
+ lv: R.Tensor((5, 10), dtype="int64") = R.one_hot(
+ indices, R.prim_value(1), R.prim_value(0), depth=10,
axis=-1
)
- lv1: R.Tensor((5, 1), dtype="int64") = R.expand_dims(indices,
axis=[-1])
- lv2: R.Tensor((5, 10), dtype="bool") = R.equal(lv1, lv)
- lv3: R.Tensor((5, 10), dtype="int64") = R.astype(lv2,
dtype="int64")
- gv: R.Tuple(R.Tensor((5, 10), dtype="int64")) = (lv3,)
+ gv: R.Tuple(R.Tensor((5, 10), dtype="int64")) = (lv,)
R.output(gv)
return gv
- example_args = (torch.randint(0, 10, (5,), dtype=torch.int64),)
+ example_args = (torch.tensor([0, 1, 5, 8, 9], dtype=torch.int64),)
- verify_model(OneHot(), example_args, {}, Expected)
+ verify_model(OneHot(), example_args, {}, Expected,
run_ep_decomposition=False)
+ verify_model_numerically(OneHot(), example_args)
def test_one_hot_invalid_num_classes():
@@ -7362,96 +7012,41 @@ def test_unflatten():
verify_model(Unflatten1(), example_args, {}, Expected)
-def test_unflatten_zero_sized_dim():
- class Unflatten(Module):
- def forward(self, x):
- return x.unflatten(0, (2, -1))
-
- verify_model_numerically(Unflatten(), (torch.randn(2, 0,
dtype=torch.float32),))
-
-
def test_gather():
- class Gather0(Module):
- def forward(self, data, indices):
- return torch.gather(data, 0, indices)
-
- class Gather1(Module):
- def forward(self, data, indices):
- return torch.gather(data, 1, indices)
-
- class Gather2(Module):
- def forward(self, data, indices):
- return torch.gather(data, -1, indices)
+ class Gather(Module):
+ def __init__(self, axis):
+ super().__init__()
+ self.axis = axis
- class Gather3(Module):
def forward(self, data, indices):
- return torch.gather(data, -2, indices)
-
- @tvm.script.ir_module
- class Expected0:
- @R.function
- def main(
- inp_0: R.Tensor((2, 3), dtype="float32"),
- inp_1: R.Tensor((2, 3), dtype="int64"),
- ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((2, 3), dtype="float32") =
R.gather_elements(inp_0, inp_1, axis=0)
- gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- @tvm.script.ir_module
- class Expected1:
- @R.function
- def main(
- inp_0: R.Tensor((2, 3), dtype="float32"),
- inp_1: R.Tensor((2, 3), dtype="int64"),
- ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((2, 3), dtype="float32") =
R.gather_elements(inp_0, inp_1, axis=1)
- gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- @tvm.script.ir_module
- class Expected2:
- @R.function
- def main(
- inp_0: R.Tensor((2, 3), dtype="float32"),
- inp_1: R.Tensor((2, 3), dtype="int64"),
- ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((2, 3), dtype="float32") =
R.gather_elements(inp_0, inp_1, axis=-1)
- gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- @tvm.script.ir_module
- class Expected3:
- @R.function
- def main(
- inp_0: R.Tensor((2, 3), dtype="float32"),
- inp_1: R.Tensor((2, 3), dtype="int64"),
- ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((2, 3), dtype="float32") =
R.gather_elements(inp_0, inp_1, axis=-2)
- gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
+ return torch.gather(data, self.axis, indices)
example_args = (
torch.randn(2, 3, dtype=torch.float32),
torch.randint(0, 3, (2, 3), dtype=torch.int64),
)
-
- verify_model(Gather0(), example_args, {}, Expected0)
- verify_model(Gather1(), example_args, {}, Expected1)
- verify_model(Gather2(), example_args, {}, Expected2)
- verify_model(Gather3(), example_args, {}, Expected3)
+ for axis in (0, 1, -1, -2):
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(
+ inp_0: R.Tensor((2, 3), dtype="float32"),
+ inp_1: R.Tensor((2, 3), dtype="int64"),
+ ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
+ with R.dataflow():
+ lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(
+ inp_0, inp_1, axis=axis
+ )
+ gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
+ R.output(gv)
+ return gv
+
+ verify_model(Gather(axis), example_args, {}, Expected)
def test_index_put():
- # Test case 1: 1D input
+ # 1D input
class IndexPut1D(Module):
def forward(self, data, indices_0, values):
indices_tuple = (indices_0,)
@@ -7479,7 +7074,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 2: 2D input
+ # 2D input
class IndexPut2D(Module):
def forward(self, data, indices_0, indices_1, values):
indices_tuple = (indices_0, indices_1)
@@ -7509,76 +7104,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 3: 3D input
- class IndexPut3D(Module):
- def forward(self, data, indices_0, indices_1, indices_2, values):
- indices_tuple = (indices_0, indices_1, indices_2)
- return data.index_put_(indices_tuple, values, accumulate=False)
-
- example_args_3d = (
- torch.randn(16, 32, 64, dtype=torch.float32),
- torch.randint(0, 16, (128,), dtype=torch.int64),
- torch.randint(0, 32, (128,), dtype=torch.int64),
- torch.randint(0, 64, (128,), dtype=torch.int64),
- torch.randn(128, dtype=torch.float32),
- )
-
- @I.ir_module
- class Expected3D:
- @R.function
- def main(
- data: R.Tensor((16, 32, 64), dtype="float32"),
- indices_0: R.Tensor((128,), dtype="int64"),
- indices_1: R.Tensor((128,), dtype="int64"),
- indices_2: R.Tensor((128,), dtype="int64"),
- values: R.Tensor((128,), dtype="float32"),
- ) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put(
- data, (indices_0, indices_1, indices_2), values,
accumulate=False
- )
- gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- # Test case 4: 4D input
- class IndexPut4D(Module):
- def forward(self, data, indices_0, indices_1, indices_2, indices_3,
values):
- indices_tuple = (indices_0, indices_1, indices_2, indices_3)
- return data.index_put_(indices_tuple, values, accumulate=False)
-
- example_args_4d = (
- torch.randn(8, 16, 32, 64, dtype=torch.float32),
- torch.randint(0, 8, (128,), dtype=torch.int64),
- torch.randint(0, 16, (128,), dtype=torch.int64),
- torch.randint(0, 32, (128,), dtype=torch.int64),
- torch.randint(0, 64, (128,), dtype=torch.int64),
- torch.randn(128, dtype=torch.float32),
- )
-
- @I.ir_module
- class Expected4D:
- @R.function
- def main(
- data: R.Tensor((8, 16, 32, 64), dtype="float32"),
- indices_0: R.Tensor((128,), dtype="int64"),
- indices_1: R.Tensor((128,), dtype="int64"),
- indices_2: R.Tensor((128,), dtype="int64"),
- indices_3: R.Tensor((128,), dtype="int64"),
- values: R.Tensor((128,), dtype="float32"),
- ) -> R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")):
- with R.dataflow():
- lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put(
- data,
- (indices_0, indices_1, indices_2, indices_3),
- values,
- accumulate=False,
- )
- gv: R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- # Test case 5: 5D input
+ # 5D input
class IndexPut5D(Module):
def forward(self, data, indices_0, indices_1, indices_2, indices_3,
indices_4, values):
indices_tuple = (indices_0, indices_1, indices_2, indices_3,
indices_4)
@@ -7617,7 +7143,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 6: 2D input with multi-dimensional index (broadcasting)
+ # 2D input with multi-dimensional index (broadcasting)
# This tests the multi-dimensional index support with broadcasting
class IndexPutBroadcast1D(Module):
def forward(self, data, indices_1):
@@ -7652,7 +7178,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 7: 2D input with multi-dimensional index (second position)
+ # 2D input with multi-dimensional index (second position)
class IndexPutBroadcast2D(Module):
def forward(self, data, indices_0):
indices_1 = torch.arange(data.shape[1]).unsqueeze(1)
@@ -7686,7 +7212,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 8: 3D input with mixed 1D and 2D indices
+ # 3D input with mixed 1D and 2D indices
class IndexPutBroadcast3D(Module):
def forward(self, data, indices_1):
indices_0 = torch.arange(data.shape[0]).unsqueeze(1)
@@ -7725,7 +7251,7 @@ def test_index_put():
R.output(gv)
return gv
- # Test case 9: batched indexing with slice (e.g., M[:, rows, cols] = x)
+ # batched indexing with slice (e.g., M[:, rows, cols] = x)
class IndexPutBatchedWithNone(Module):
def forward(self, x):
B = x.size(0)
@@ -7765,8 +7291,6 @@ def test_index_put():
# Run verification for each case
verify_model(IndexPut1D(), example_args_1d, {}, Expected1D)
verify_model(IndexPut2D(), example_args_2d, {}, Expected2D)
- verify_model(IndexPut3D(), example_args_3d, {}, Expected3D)
- verify_model(IndexPut4D(), example_args_4d, {}, Expected4D)
verify_model(IndexPut5D(), example_args_5d, {}, Expected5D)
verify_model(IndexPutBroadcast1D(), example_args_broadcast1, {},
ExpectedBroadcast1D)
verify_model(IndexPutBroadcast2D(), example_args_broadcast2, {},
ExpectedBroadcast2D)
@@ -7801,40 +7325,6 @@ def test_index_put_with_tuple_output():
)
-def test_m4d_diag_index_put_tuple_output_regression():
- class M4D(Module):
- def forward(self, x):
- b, k, n = 2, 3, 5
- buf = x.new_zeros(b, k, n, n)
- idx = torch.arange(n, device=x.device)
-
- diag = buf[..., idx, idx]
- diag = torch.nn.functional.elu(diag) + 1.0 + 1e-8
- buf[..., idx, idx] = diag
-
- return x[..., :1], buf
-
- ex_in = torch.zeros(2, 3, 5, dtype=torch.float32)
- exported_program = export(M4D().eval(), args=(ex_in,))
-
- exported_targets = [str(getattr(n, "target", "")) for n in
exported_program.graph.nodes]
- assert any("index_put" in target for target in exported_targets)
-
- # Regression focus: importing this graph should not segfault at Tuple
construction.
- mod = from_exported_program(exported_program)
- ret_ty = mod["main"].ret_ty
- assert isinstance(ret_ty, relax.TupleType)
-
- tensor_fields = [f for f in ret_ty.fields if isinstance(f,
relax.TensorType)]
- assert len(tensor_fields) >= 2
- # x: (2, 3, 5) → x[..., :1]: (2, 3, 1)
- assert any(len(f.shape) == 3 and int(f.shape[-1]) == 1 for f in
tensor_fields)
- # buf: (2, 3, 5, 5) → 4-D with spatial dims 5x5
- assert any(
- len(f.shape) == 4 and int(f.shape[-2]) == 5 and int(f.shape[-1]) == 5
for f in tensor_fields
- )
-
-
def test_index_put_mutation_through_alias_regression():
class IndexPutAlias(Module):
def forward(self, x, idx, values):
@@ -8959,39 +8449,6 @@ def test_dynamic_shape_with_addition_constraints():
)
-def test_dynamic_shape_with_subtraction_constraints():
- class ConcatModel(torch.nn.Module):
- def forward(self, x, y):
- return torch.cat([x, y], dim=0)
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- 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},
- "tir_var_upper_bound": {"s17": 63},
- }
- )
- with R.dataflow():
- 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
-
- batch = torch.export.Dim("batch", min=1, max=64)
- example_args = (torch.randn(8, 4), torch.randn(7, 4))
- dynamic_shapes = {"x": {0: batch}, "y": {0: batch - 1}}
-
- verify_model(
- ConcatModel(), example_args, {}, Expected,
dynamic_shapes=dynamic_shapes, map_free_vars=True
- )
-
-
def test_dynamic_shape_with_multiplication_constraints():
class ConcatModel(torch.nn.Module):
def forward(self, x, y):
@@ -9345,57 +8802,6 @@ def test_torchvision_roi_align_aligned():
verify_model_numerically(ROIAlign(), example_args, rtol=1e-5, atol=1e-5)
-def test_upsample_nearest2d():
- class UpsampleNearest2dScale(Module):
- def forward(self, input):
- return torch.nn.functional.interpolate(input, scale_factor=2.0,
mode="nearest")
-
- class UpsampleNearest2dSize(Module):
- def forward(self, input):
- return torch.nn.functional.interpolate(input, size=(20, 20),
mode="nearest")
-
- example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
-
- @tvm.script.ir_module
- class expected_scale:
- @R.function
- def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) ->
R.Tuple(
- R.Tensor((1, 3, 20, 20), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 20, 20), dtype="float32") =
R.image.resize2d(
- input_1,
- size=(20, 20),
- layout="NCHW",
- method="nearest_neighbor",
- coordinate_transformation_mode="half_pixel",
- )
- gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- @tvm.script.ir_module
- class expected_size:
- @R.function
- def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) ->
R.Tuple(
- R.Tensor((1, 3, 20, 20), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 20, 20), dtype="float32") =
R.image.resize2d(
- input_1,
- size=(20, 20),
- layout="NCHW",
- method="nearest_neighbor",
- coordinate_transformation_mode="half_pixel",
- )
- gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,)
- R.output(gv)
- return gv
-
- verify_model(UpsampleNearest2dScale(), example_args, {}, expected_scale)
- verify_model(UpsampleNearest2dSize(), example_args, {}, expected_size)
-
-
def test_from_exported_program_sparse_csr_buffer():
class SparseCsrBufferModule(nn.Module):
def __init__(self):
@@ -9421,6 +8827,48 @@ def test_from_exported_program_sparse_csr_buffer():
assert isinstance(mod, tvm.IRModule)
[email protected](not env.has_llvm(), reason="need llvm")
[email protected]("with_message,in_cond", [(False, False), (True,
False), (True, True)])
+def test_assert_async(with_message, in_cond):
+ class AssertAsync(Module):
+ def forward(self, x):
+ def checked(value):
+ if with_message:
+ torch.ops.aten._assert_async.msg(
+ value.clamp(min=0), assert_msg="Positive required:
{value}"
+ )
+ else:
+ torch.ops.aten._assert_async.default(value)
+ return value.clone()
+
+ if in_cond:
+ return torch.cond(x >= 0, checked, lambda value:
value.clone(), (x,))
+ return checked(x)
+
+ # Also cover non-boolean conditions and one-element, non-scalar tensors.
+ example_args = (torch.tensor(0.5) if with_message else
torch.tensor([True]),)
+ model = AssertAsync()
+ # PyTorch's decomposition removes the message-less overload.
+ mod = from_exported_program(export(model, args=example_args),
run_ep_decomposition=with_message)
+ assert not mod["main"].is_pure
+ vm = relax.VirtualMachine(relax.build(mod, target="llvm"), tvm.cpu())
+ actual = vm["main"](tvm.runtime.tensor(example_args[0].numpy()))
+ np.testing.assert_array_equal(actual[0].numpy(),
model(*example_args).numpy())
+
+ invalid = torch.zeros_like(example_args[0])
+ with pytest.raises(RuntimeError):
+ model(invalid)
+ message = "Positive required: [{]value[}]" if with_message else "Assertion
Failed"
+ with pytest.raises(AssertionError, match=message):
+ vm["main"](tvm.runtime.tensor(invalid.numpy()))
+
+ if in_cond:
+ # The unchecked branch must not execute the other branch's assertion.
+ negative = -example_args[0]
+ actual = vm["main"](tvm.runtime.tensor(negative.numpy()))
+ np.testing.assert_array_equal(actual[0].numpy(),
model(negative).numpy())
+
+
def test_cond_basic():
"""Basic data-dependent cond with runtime predicate."""
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index b7074a5f94..8002a03c25 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -710,9 +710,6 @@ def test_linear():
# matmul
class MatMul1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return torch.matmul(x, y)
@@ -742,9 +739,6 @@ def test_linear():
def test_bmm():
class BMM(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return torch.bmm(x, y)
@@ -774,16 +768,10 @@ def test_bmm():
def test_baddbmm():
class BAddBMM1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, c, x, y):
return torch.baddbmm(c, x, y)
class BAddBMM2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, c, x, y):
return torch.baddbmm(c, x, y, alpha=2, beta=0)
@@ -836,16 +824,10 @@ def test_baddbmm():
def test_einsum():
class Einsum1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x):
return torch.einsum("ii", x)
class Einsum2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return torch.einsum("i,j->ij", x, y)
@@ -1008,9 +990,6 @@ def test_maxpool1d():
return self.pool(input)
class MaxPool1d_functional(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool1d(input, kernel_size=2)
@@ -1108,9 +1087,6 @@ def test_maxpool2d():
return self.pool(input)
class MaxPool2d_functional(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1])
@@ -1211,9 +1187,6 @@ def test_maxpool3d():
return self.pool(input)
class MaxPool3d_functional(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, input):
return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1])
@@ -2017,35 +1990,12 @@ def test_functional_layernorm():
def forward(self, input):
return torch.nn.functional.layer_norm(input, self.shape,
self.weight, self.bias, 1e-5)
- @tvm.script.ir_module
- class expected3:
- @R.function
- def main(
- input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
- w1: R.Tensor([10, 10], dtype="float32"),
- w2: R.Tensor([10, 10], dtype="float32"),
- ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
- # block 0
- with R.dataflow():
- lv: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.nn.layer_norm(
- input_1,
- w1,
- w2,
- axes=[-2, -1],
- epsilon=1e-05,
- center=True,
- scale=True,
- )
- gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
- R.output(gv)
- return gv
-
model = LayerNorm3([10, 10])
binding = {
"w1": model.weight.detach().numpy(),
"w2": model.bias.detach().numpy(),
}
- verify_model(model, input_info, binding, expected3)
+ verify_model(model, input_info, binding, expected1)
def test_cross_entropy():
@@ -2556,112 +2506,6 @@ def test_div_mode():
verify_model(DivFloorModel(), input_info, {}, expected_div_floor)
-def test_round_decimals():
- """torch.round(x, decimals) through from_fx must match PyTorch's
round-half-to-even
- results, including negative decimals (round(25, -1) == 20). The previous
- scale-by-10**decimals implementation multiplied by 0.1 for negative
decimals, which
- is numerically wrong: 25 * 0.1 == 2.5000000000000004 in float64 rounds up
to 30.
- """
- input_info = [([10], "float32")]
- x = torch.tensor(
- [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25],
dtype=torch.float32
- )
-
- class RoundDecimalsModel(Module):
- def __init__(self, decimals):
- super().__init__()
- self.decimals = decimals
-
- def forward(self, input):
- return torch.round(input, decimals=self.decimals)
-
- for decimals in (0, 1, -1, -2):
- gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
- mod = from_fx(gm, input_info)
- ex = relax.build(mod, target="llvm")
- vm = relax.VirtualMachine(ex, tvm.cpu())
- tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
- got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
- tvm.testing.assert_allclose(
- got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6,
atol=1e-6
- )
-
-
-def test_round_decimals_low_precision():
- """Scaling for low-precision inputs must happen in float32 and be cast
back.
-
- 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 ==
250000
- exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5
- already overflows float16 (the scale itself becomes inf), turning
decimals=5
- and -5 into NaN. Upcasting the input to float32 keeps the scaling exact;
the
- rounded result is cast back to the input dtype.
- """
- input_info = [([8], "float16")]
- x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5],
dtype=torch.float16)
-
- class RoundDecimalsModel(Module):
- def __init__(self, decimals):
- super().__init__()
- self.decimals = decimals
-
- def forward(self, input):
- return torch.round(input, decimals=self.decimals)
-
- # Positive decimals exercise the multiply-by-10**d overflow (4, 5);
- # negative decimals exercise the 10**|d| scale overflowing float16 (-5).
- for decimals in (2, 4, 5, -2, -4, -5):
- gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
- mod = from_fx(gm, input_info)
- ex = relax.build(mod, target="llvm")
- vm = relax.VirtualMachine(ex, tvm.cpu())
- tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
- got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
- tvm.testing.assert_allclose(
- got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6,
atol=1e-6
- )
-
-
-def test_round_decimals_large():
- """A large |decimals| must import and run without OverflowError.
-
- The scale 10**|decimals| used to be built as an unbounded host Python int
- before being handed to relax.const, whose int-to-float conversion raises
- OverflowError ("int too large to convert to float") once |decimals| >= 309
- (10**309 already exceeds the float64 range). PyTorch accepts such decimals
--
- torch.round(x, decimals=309) -- and traces a valid round.decimals call, so
- importing the graph must not crash on them. The scale is now built directly
- in the float dtype and saturates to inf once it leaves the finite range,
- matching PyTorch, whose all-NaN result here comes from the same inf scale.
- """
- input_info = [([5], "float32")]
- x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32)
-
- class RoundDecimalsModel(Module):
- def __init__(self, decimals):
- super().__init__()
- self.decimals = decimals
-
- def forward(self, input):
- return torch.round(input, decimals=self.decimals)
-
- for decimals in (309, -309):
- gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
- mod = from_fx(gm, input_info) # used to raise OverflowError here
- ex = relax.build(mod, target="llvm")
- vm = relax.VirtualMachine(ex, tvm.cpu())
- tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
- got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
-
- # The scale overflows to inf, and IEEE arithmetic turns every element
into
- # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
- # (empty here) finite elements separately, since allclose fails on NaN.
- expected = torch.round(x, decimals=decimals)
- actual = torch.as_tensor(got)
- assert torch.equal(torch.isnan(actual), torch.isnan(expected))
- finite = ~torch.isnan(expected)
- assert torch.allclose(actual[finite], expected[finite], rtol=1e-6,
atol=1e-6)
-
-
def test_size():
input_info = [([1, 3, 10, 10], "float32")]
@@ -3930,30 +3774,7 @@ def test_interpolate():
align_corners=False,
)
- @tvm.script.ir_module
- class expected7:
- @R.function
- def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) ->
R.Tensor(
- (1, 3, 8, 40, 40), dtype="float32"
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") =
R.image.resize3d(
- input_5,
- (8, 40, 40),
- roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000,
0.000000],
- layout="NCDHW",
- method="linear",
- coordinate_transformation_mode="half_pixel",
- rounding_method="",
- cubic_alpha=-0.75,
- cubic_exclude=0,
- extrapolation_value=0,
- )
- gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
- R.output(gv)
- return gv
-
- verify_model(Interpolate7(), input_info_5d, {}, expected7)
+ verify_model(Interpolate7(), input_info_5d, {}, expected6)
class Interpolate8(Module):
def forward(self, input):
@@ -4150,16 +3971,10 @@ def test_addmm():
]
class Addmm1(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3)
class Addmm2(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
@@ -4610,19 +4425,7 @@ def test_datatype():
def forward(self, x):
return x.type(torch.float32)
- # type
- class TypeFromAttr(Module):
- def forward(self, x):
- return x.type(x.getattr("dtype"))
-
- # astype
- class AsType(Module):
- def forward(self, x):
- return x.astype(torch.float32)
-
verify_model(Type(), input_info, {}, expected1)
- verify_model(TypeFromAttr(), input_info, {}, expected1)
- verify_model(AsType(), input_info, {}, expected1)
def test_meshgrid():
@@ -5031,9 +4834,6 @@ def test_keep_params():
def test_unwrap_unit_return_tuple():
class Identity(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x):
return (x,)
@@ -5055,9 +4855,6 @@ def test_unwrap_unit_return_tuple():
def test_no_bind_return_tuple():
class Identity(Module):
- def __init__(self):
- super().__init__()
-
def forward(self, x, y):
return (x, y)
@@ -5083,16 +4880,10 @@ def test_no_bind_return_tuple():
def test_argmax():
class Argmax1(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmax(input, dim=-1)
class Argmax2(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmax(input, dim=-1, keepdim=True)
@@ -5122,16 +4913,10 @@ def test_argmax():
def test_argmin():
class Argmin1(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmin(input)
class Argmin2(Module):
- def __init__(self) -> None:
- super().__init__()
-
def forward(self, input):
return torch.argmin(input, keepdim=True)
@@ -5629,11 +5414,7 @@ def test_slice_scatter():
def test_masked_scatter():
- class MaskedScatter1(Module):
- def forward(self, data, mask, src):
- return data.masked_scatter(mask, src)
-
- class MaskedScatter2(Module):
+ class MaskedScatter(Module):
def forward(self, data, mask, src):
return data.masked_scatter(mask, src)
@@ -5679,10 +5460,10 @@ def test_masked_scatter():
return gv
verify_model(
- MaskedScatter1(), [([5], "float32"), ([5], "bool"), ([10],
"float32")], {}, expected1
+ MaskedScatter(), [([5], "float32"), ([5], "bool"), ([10], "float32")],
{}, expected1
)
verify_model(
- MaskedScatter2(),
+ MaskedScatter(),
[([2, 5], "float32"), ([2, 5], "bool"), ([3, 5], "float32")],
{},
expected2,
@@ -5772,27 +5553,6 @@ def test_flip():
verify_model(Flip0(), [([2, 2], "float32")], {}, Expected0)
-def test_flip_multi_axis():
- class FlipMulti(Module):
- def forward(self, data):
- return torch.flip(data, [0, 1])
-
- @tvm.script.ir_module
- class ExpectedMulti:
- @R.function
- def main(
- inp_0: R.Tensor((2, 3), dtype="float32"),
- ) -> R.Tensor((2, 3), dtype="float32"):
- with R.dataflow():
- lv: R.Tensor((2, 3), dtype="float32") = R.flip(inp_0, axis=0)
- lv1: R.Tensor((2, 3), dtype="float32") = R.flip(lv, axis=1)
- gv: R.Tensor((2, 3), dtype="float32") = lv1
- R.output(gv)
- return gv
-
- verify_model(FlipMulti(), [([2, 3], "float32")], {}, ExpectedMulti)
-
-
def test_take():
class Take(Module):
def forward(self, data, indices):
@@ -6027,7 +5787,7 @@ def test_select():
def test_inplace_copy():
- class Inplace_Copy(Module):
+ class Copy(Module):
def forward(self, x, y):
x.copy_(y)
return x
@@ -6047,11 +5807,6 @@ def test_inplace_copy():
R.output(gv)
return gv
- class CopyBroadcast(Module):
- def forward(self, x, src):
- x.copy_(src)
- return x
-
@tvm.script.ir_module
class expected_copy:
@R.function
@@ -6066,12 +5821,12 @@ def test_inplace_copy():
return gv
verify_model(
- Inplace_Copy(),
+ Copy(),
[((1, 2, 3, 4), "float32"), ((1, 2, 3, 4), "float32")],
{},
Expected,
)
- verify_model(CopyBroadcast(), [((2, 3), "float32"), ((), "int64")], {},
expected_copy)
+ verify_model(Copy(), [((2, 3), "float32"), ((), "int64")], {},
expected_copy)
def test_clone():
@@ -6537,12 +6292,12 @@ def test_round():
input_info = [([3, 4], "float32")]
class Round(Module):
- def __init__(self, decimals=0):
+ def __init__(self, decimals=None):
super().__init__()
self.decimals = decimals
def forward(self, x):
- if self.decimals == 0:
+ if self.decimals is None:
return torch.round(x)
else:
return torch.round(x, decimals=self.decimals)
@@ -6574,38 +6329,29 @@ def test_round():
return gv
rounds = [
- (0, Expected1),
+ (None, Expected1),
(2, Expected2),
]
for decimals, expected in rounds:
verify_model(Round(decimals), input_info, {}, expected)
- # Test numerical accuracy with decimals
- test_data = torch.tensor(
- [
- [1.2345, 2.3456, 3.4567, 4.5678],
- [5.6789, 6.7890, 7.8901, 8.9012],
- [9.1234, 10.2345, 11.3456, 12.4567],
- ]
- )
-
- for decimals in [0, 2]:
- torch_model = Round(decimals)
- graph_model = fx.symbolic_trace(torch_model)
- with torch.no_grad():
- mod = from_fx(graph_model, input_info)
-
- target = tvm.target.Target("llvm")
- ex = relax.build(mod, target)
- vm = relax.VirtualMachine(ex, tvm.cpu())
-
- torch_result = torch_model(test_data).numpy()
- tvm_input = tvm.runtime.tensor(test_data.numpy())
- tvm_result = vm["main"](tvm_input).numpy()
-
- # Use relaxed tolerance due to floating-point precision in decimal
operations
- tvm.testing.assert_allclose(tvm_result, torch_result, rtol=1e-3,
atol=1e-3)
+ # Float16 needs float32 scaling to avoid intermediate overflow.
+ cases = [
+ (torch.float32, (0, 1, -1, -2)),
+ (torch.float16, (2, 4, 5, -2, -4, -5)),
+ ]
+ for dtype, decimals_values in cases:
+ x = torch.tensor(
+ [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25, 0.0],
dtype=dtype
+ )
+ for decimals in decimals_values:
+ model = Round(decimals).eval()
+ mod = from_fx(fx.symbolic_trace(model), [(x.shape, dtype)])
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ actual = vm["main"](tvm.runtime.tensor(x.numpy())).numpy()
+ tvm.testing.assert_allclose(actual, model(x).numpy(), rtol=1e-6,
atol=1e-6)
if __name__ == "__main__":
