| """ |
| PYTEST_DONT_REWRITE (prevents pytest from rewriting assertions, which interferes |
| with test_functionalization_with_native_python_assertion) |
| """ |
| |
| # Owner(s): ["module: dynamo"] |
| import unittest |
| from typing import List, Set |
| import operator |
| |
| import torch |
| from torch.testing._internal.common_utils import run_tests, TestCase |
| from torch.testing import FileCheck |
| from torch._dynamo.eval_frame import is_dynamo_supported |
| from torch._export import export, dynamic_dim |
| from torch._export.constraints import constrain_as_value, constrain_as_size |
| from torch._export.exported_program import ExportGraphSignature |
| from torch._export.passes import ( |
| ReplaceViewOpsWithViewCopyOpsPass, |
| ) |
| from torch._export.passes.replace_view_ops_with_view_copy_ops_pass import ( |
| is_view_op, |
| get_view_copy_of_view_op, |
| ) |
| from torch._export.passes.functionalize_side_effectful_ops_pass import ( |
| _FunctionalizeSideEffectfulOpsPass, |
| ) |
| from functorch.experimental.control_flow import cond |
| from torch.fx.passes.operator_support import OperatorSupport |
| from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition |
| from torch.fx._symbolic_trace import symbolic_trace |
| from torch.utils._pytree import tree_flatten |
| |
| |
| def count_call_function(graph: torch.fx.Graph, target: torch.ops.OpOverload) -> int: |
| count = 0 |
| for node in graph.nodes: |
| if node.op == "call_function" and node.target == target: |
| count += 1 |
| return count |
| |
| |
| class _AddOperatorSupport(OperatorSupport): |
| def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: |
| return node.op == "call_function" and node.target in {operator.add} |
| |
| |
| class _AtenAddOperatorSupport(OperatorSupport): |
| def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: |
| return node.op == "call_function" and node.target in { |
| torch.ops.aten.add.Tensor |
| } |
| |
| |
| def _to_partition_names(partitions: List[Partition]) -> List[Set[str]]: |
| return [{n.name for n in p.nodes} for p in partitions] |
| |
| |
| def _get_output_names(gm: torch.fx.GraphModule) -> List[str]: |
| output_node = next(n for n in gm.graph.nodes if n.op == "output") |
| args = tree_flatten(output_node.args)[0] |
| # if isinstance(args, tuple) and len(args) == 1: |
| # args = args[0] |
| return [str(arg) for arg in args] |
| |
| |
| @unittest.skipIf(not is_dynamo_supported(), "Dynamo not supported") |
| class TestPasses(TestCase): |
| def test_replace_broken_ops(self) -> None: |
| x = torch.randn([2, 3, 4, 5]) |
| model: torch.nn.Linear = torch.nn.Linear(5, 5) |
| |
| def f(inp: torch.Tensor) -> torch.Tensor: |
| return model(inp) |
| |
| ep = export(f, (x,)).transform(ReplaceViewOpsWithViewCopyOpsPass()) |
| |
| count_after = 0 |
| for node in ep.graph.nodes: |
| if node.target == torch.ops.aten.view.default: |
| count_after += 1 |
| self.assertEqual(count_after, 0) |
| self.assertTrue(torch.allclose(ep(x), f(x))) |
| |
| def test_runtime_assert_one_dim(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x): |
| return x.cos() |
| |
| x = torch.zeros(2, 2, 3) |
| |
| ep = export(M(), (x,), constraints=[dynamic_dim(x, 1) >= 2, dynamic_dim(x, 1) <= 6]) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| |
| self.assertEqual(num_assert, 3) |
| self.assertEqual(num_scalar_tensor, 3) |
| |
| with self.assertRaisesRegex(RuntimeError, "Input arg0_1"): |
| ep(torch.zeros(2, 7, 3)) |
| |
| self.assertEqual(ep(torch.ones(2, 4, 3)), M().forward(torch.ones(2, 4, 3))) |
| |
| def test_runtime_assert_multiple_dims(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x, y): |
| return x.cos().sum() + y.sin().sum() |
| |
| x = torch.zeros(4, 2, 3) |
| y = torch.zeros(5, 5, 5) |
| |
| constraints = [ |
| dynamic_dim(x, 1) >= 2, |
| dynamic_dim(x, 1) <= 6, |
| dynamic_dim(y, 0) >= 3, |
| dynamic_dim(x, 0) >= 3 |
| ] |
| |
| ep = export(M(), (x, y), constraints=constraints) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| |
| self.assertEqual(num_assert, 6) |
| self.assertEqual(num_scalar_tensor, 6) |
| |
| with self.assertRaisesRegex(RuntimeError, "Input arg0_1"): |
| ep(torch.zeros(4, 7, 3), torch.ones(5, 5, 5)) |
| |
| with self.assertRaisesRegex(RuntimeError, "Input arg1_1"): |
| ep(torch.zeros(4, 2, 3), torch.ones(2, 5, 5)) |
| |
| def test_runtime_assert_some_dims_not_specified(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x, y): |
| return x.cos().sum() + y.sin().sum() |
| |
| x = torch.zeros(4, 2, 3) |
| y = torch.zeros(5, 5, 5) |
| |
| constraints = [ |
| dynamic_dim(x, 1) >= 2, |
| dynamic_dim(x, 1) <= 6, |
| dynamic_dim(x, 0) >= 3 |
| ] |
| |
| ep = export(M(), (x, y), constraints=constraints) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| |
| # there are 3 asserts from y and 2 from dynamic x dims and 1 from static x dim |
| self.assertEqual(num_assert, 6) |
| self.assertEqual(num_scalar_tensor, 6) |
| |
| with self.assertRaisesRegex(RuntimeError, "Input arg0_1"): |
| ep(torch.zeros(4, 7, 3), torch.ones(5, 5, 5)) |
| |
| # y is specialized to 5 |
| with self.assertRaisesRegex(RuntimeError, r"Input arg1_1.shape\[0\] is specialized at 5"): |
| ep(torch.zeros(4, 2, 3), torch.ones(2, 5, 5)) |
| |
| # Since we didn't insert the constraint for x[1] >= 2, it should work for case where x[1] == 1 |
| gm_result_for_1_size = ep(torch.ones(3, 1, 3), torch.ones(5, 5, 5)) |
| eager_result_for_1_size = M().forward(torch.ones(3, 1, 3), torch.ones(5, 5, 5)) |
| |
| self.assertEqual(gm_result_for_1_size, eager_result_for_1_size) |
| |
| def test_runtime_assert_some_inps_not_used(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x, y): |
| return y.cos().sum() |
| |
| x = torch.zeros(4, 2, 3) |
| y = torch.zeros(5, 5, 5) |
| |
| constraints = [ |
| dynamic_dim(y, 1) >= 3, |
| dynamic_dim(y, 1) <= 6, |
| ] |
| |
| ep = export(M(), (x, y), constraints=constraints) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| |
| # there are 4 asserts from y and 3 from x |
| self.assertEqual(num_assert, 7) |
| self.assertEqual(num_scalar_tensor, 7) |
| |
| with self.assertRaisesRegex(RuntimeError, "Input arg0_1"): |
| ep(torch.zeros(4, 7, 3), torch.ones(5, 5, 5)) |
| |
| # y is specialized to 5 |
| with self.assertRaisesRegex(RuntimeError, r"Input arg1_1.shape\[0\] is specialized at 5"): |
| ep(torch.zeros(4, 2, 3), torch.ones(2, 5, 5)) |
| |
| # Since we didn't insert the constraint for x[1] >= 2, it should work for case where x[1] == 1 |
| gm_result_for_1_size = ep(torch.zeros(4, 2, 3), torch.ones(5, 5, 5)) |
| eager_result_for_1_size = M().forward(torch.zeros(4, 2, 3), torch.ones(5, 5, 5)) |
| |
| self.assertEqual(gm_result_for_1_size, eager_result_for_1_size) |
| |
| def test_view_to_view_copy(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x): |
| z = x.view(x.shape) |
| return z.cos().sum() |
| |
| x = torch.zeros(4, 2, 3) |
| |
| ep = export(M(), (x,)) |
| self.assertEqual(count_call_function(ep.graph, torch.ops.aten.view.default), 1) |
| |
| ep = ep.transform(ReplaceViewOpsWithViewCopyOpsPass()) |
| self.assertEqual(count_call_function(ep.graph, torch.ops.aten.view.default), 0) |
| |
| def test_functionalization_with_view_copy(self) -> None: |
| def foo(x): |
| y = x + 4 |
| y.add_(4) |
| z = y.view(y.shape) |
| return x.cos() + z.cos() |
| |
| x = torch.zeros(4, 2, 3) |
| |
| ep = export(foo, (x,)).transform(ReplaceViewOpsWithViewCopyOpsPass()) |
| # After this pass, there shouldn't be any view nodes in the graph |
| self.assertTrue(count_call_function(ep.graph, torch.ops.aten.view.default) == 0) |
| self.assertTrue(count_call_function(ep.graph, torch.ops.aten.view_copy.default) > 0) |
| |
| def test_views_op_having_view_copy(self) -> None: |
| schemas = torch._C._dispatch_get_registrations_for_dispatch_key("") |
| aten_schemas = [s[6:] for s in schemas if s.startswith("aten::")] |
| |
| for aten_schema in aten_schemas: |
| val = aten_schema.split(".") |
| assert len(val) <= 2 |
| name = "" |
| overload = "" |
| if len(val) == 1: |
| name = val[0] |
| overload = "default" |
| else: |
| name, overload = val[0], val[1] |
| |
| op_overload = getattr(getattr(torch.ops.aten, name), overload) |
| if torch.Tag.core in op_overload.tags and is_view_op(op_overload._schema): |
| self.assertIsNotNone(get_view_copy_of_view_op(op_overload._schema)) |
| |
| def test_runtime_assert_inline_constraints_for_item(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x): |
| b = x.item() |
| constrain_as_value(b, min=2, max=5) |
| return b |
| |
| x = torch.tensor([2]) |
| mod = M() |
| ep = export(mod, (x,)) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| # 1 constraint for shape of x, 2 constraints for b |
| self.assertEqual(num_assert, 3) |
| self.assertEqual(num_scalar_tensor, 3) |
| |
| with self.assertRaisesRegex(RuntimeError, r"_local_scalar_dense_default is outside of inline constraint \[2, 5\]."): |
| ep(torch.tensor([6])) |
| |
| new_inp = torch.tensor([5]) |
| self.assertEqual(mod(new_inp), ep(new_inp)) |
| |
| def test_runtime_assert_inline_constraints_for_nonzero(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, x): |
| b = x.nonzero() |
| constrain_as_value(b.shape[0], min=3, max=5) |
| return b |
| |
| x = torch.tensor([2, 1, 2, 3, 5, 0]) |
| |
| mod = M() |
| ep = export(mod, (x,), constraints=[dynamic_dim(x, 0) >= 2]) |
| |
| num_assert = count_call_function(ep.graph, torch.ops.aten._assert_async.msg) |
| num_scalar_tensor = count_call_function(ep.graph, torch.ops.aten.scalar_tensor.default) |
| |
| # TODO: De-duplicate assertions for same symbol. |
| self.assertEqual(num_assert, 4) |
| self.assertEqual(num_scalar_tensor, 4) |
| |
| with self.assertRaisesRegex(RuntimeError, r"nonzero_default.shape\[0\] is outside of inline constraint \[3, 5\]."): |
| ep(torch.tensor([1, 1, 0, 0, 0])) |
| |
| with self.assertRaisesRegex(RuntimeError, r"nonzero_default.shape\[0\] is outside of inline constraint \[3, 5\]."): |
| ep(torch.ones(6)) |
| |
| new_inp = torch.tensor([1, 1, 1, 1]) |
| self.assertEqual(mod(new_inp), ep(new_inp)) |
| |
| def test_runtime_assert_inline_constraints_for_cond(self) -> None: |
| class M(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| def forward(self, pred, x, y): |
| def true_fn(x, y): |
| b = x.item() |
| constrain_as_value(b, min=2, max=5) |
| return x - b |
| |
| def false_fn(x, y): |
| c = y.item() |
| constrain_as_value(c, min=2, max=5) |
| return y - c |
| |
| ret = cond(pred, true_fn, false_fn, [x, y]) |
| return ret |
| |
| x = torch.tensor([2]) |
| y = torch.tensor([5]) |
| mod = M() |
| ep = export(mod, (torch.tensor(True), x, y)) |
| with self.assertRaisesRegex(RuntimeError, "is outside of inline constraint \\[2, 5\\]."): |
| ep(torch.tensor(False), torch.tensor([6]), torch.tensor([6])) |
| |
| def test_runtime_assert_equality_constraint(self): |
| class Adder(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| |
| def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| return x + y |
| |
| m = Adder() |
| x = torch.rand(3, 4) |
| y = torch.rand(3, 4) |
| exported = torch._export.export( |
| m, (x, y), constraints=[dynamic_dim(x, 1) == dynamic_dim(y, 1)] |
| ) |
| |
| x = torch.rand(3, 5) |
| y = torch.rand(3, 6) |
| with self.assertRaisesRegex( |
| RuntimeError, |
| r"Input arg0_1.shape\[1\] is not equal to input arg1_1.shape\[1\]" |
| ): |
| exported(x, y) |
| |
| y = torch.rand(3, 5) |
| dynamo_result = exported(x, y) |
| real_result = m(x, y) |
| self.assertTrue(torch._dynamo.utils.same(real_result, dynamo_result)) |
| |
| def test_functionalize_inline_contraints(self) -> None: |
| def f(x): |
| a = x.item() |
| constrain_as_size(a, 4, 7) |
| return torch.empty((a, 4)) |
| |
| ep = torch._export.export(f, (torch.tensor([7]),)) |
| gm = ep.graph_module |
| FileCheck().check_count( |
| "torch.ops.aten.sym_constrain_range.default", |
| 1, |
| exactly=True, |
| ).run(gm.code) |
| |
| gm = ep.transform(_FunctionalizeSideEffectfulOpsPass()).graph_module |
| |
| with self.assertRaisesRegex( |
| RuntimeError, |
| r"_local_scalar_dense_default is outside of inline constraint \[4, 7\]", |
| ) as cm: |
| gm(torch.tensor([20])) |
| |
| inp = torch.tensor([5]) |
| res, dep_token = gm(inp) |
| self.assertEqual(res.shape, torch.Size([5, 4])) |
| self.assertEqual(dep_token.shape, torch.Size([])) |
| |
| FileCheck().check_count( |
| "torch.ops.aten._functional_sym_constrain_range", 1, exactly=True |
| ).run(gm.code) |
| FileCheck().check_count( |
| "torch.ops.aten.sym_constrain_range.default", 0, exactly=True |
| ).run(gm.code) |
| |
| dep_token_node = next(n for n in gm.graph.nodes if n.name == "dep_token3") |
| constrain_node = next( |
| n |
| for n in gm.graph.nodes |
| if n.target == torch.ops.aten._functional_sym_constrain_range |
| ) |
| self.assertEqual(constrain_node.kwargs["dep_token"], dep_token_node) |
| |
| def test_functionalize_input_constraints(self) -> None: |
| def f(x): |
| return x * 2 |
| |
| inp = torch.zeros(4, 8) |
| ep = torch._export.export( |
| f, |
| (inp,), |
| constraints=[ |
| dynamic_dim(inp, 0) < 10, |
| dynamic_dim(inp, 0) >= 3, |
| ], |
| ) |
| FileCheck().check_count( |
| "torch.ops.aten._assert_async.msg", 3, exactly=True |
| ).run(ep.graph_module.code) |
| |
| gm = ep.transform(_FunctionalizeSideEffectfulOpsPass()).graph_module |
| with self.assertRaisesRegex( |
| RuntimeError, |
| r"Input arg0_1.shape\[0\] is outside of specified dynamic range \[3, 9\]", |
| ): |
| gm(torch.ones(11, 8)) |
| |
| inp = torch.ones(6, 8) |
| self.assertEqual(gm(inp)[0], f(inp)) |
| FileCheck().check_count( |
| "torch.ops.aten._functional_assert_async.msg", 3, exactly=True |
| ).run(gm.code) |
| FileCheck().check_count( |
| "torch.ops.aten._assert_async.msg", 0, exactly=True |
| ).run(gm.code) |
| |
| def test_functionalization(self) -> None: |
| def f(x, y): |
| a = x.item() |
| constrain_as_size(a, 4, 7) |
| return x + 4, x + y * 2 |
| |
| inps = (torch.tensor([5]), torch.zeros((3, 4))) |
| ep = torch._export.export( |
| f, |
| inps, |
| constraints=[dynamic_dim(inps[1], 1) < 6], |
| _functionalize_runtime_assertions=True, |
| ) |
| FileCheck().check_count( |
| "torch.ops.aten._functional_sym_constrain_range", 1, exactly=True |
| ).run(ep.graph_module.code) |
| inps = (torch.tensor([7]), torch.ones((3, 5))) |
| self.assertTrue(torch._dynamo.utils.same(ep(*inps), f(*inps))) |
| |
| def test_functionalization_with_native_python_assertion(self) -> None: |
| def f(x): |
| b = x.sin() |
| assert x[0] == 3 |
| return x.cos() + b |
| |
| inp = torch.Tensor([3, 4, 5]) |
| ep = torch._export.export(f, (inp,), _functionalize_runtime_assertions=True) |
| |
| # Check native assertion has corresponding functional assertion nodes generated. |
| select_int_node = next( |
| n |
| for n in ep.graph_module.graph.nodes |
| if n.target == torch.ops.aten.select.int |
| ) |
| equal_scalar_node = select_int_node.next |
| dep_token_node = next( |
| n |
| for n in ep.graph_module.graph.nodes |
| if ( |
| n.target == torch.ops.aten._functional_assert_async.msg |
| and n.args[0] == equal_scalar_node |
| ) |
| ) |
| self.assertIn( |
| "call_function[target=torch.ops.aten._functional_assert_async.msg]" |
| "(args = (%eq_scalar, assertion error), kwargs = {dep_token: %dep_token1}", |
| dep_token_node.format_node(), |
| ) |
| |
| def test_functionalization_with_mutated_buffer(self) -> None: |
| buf = torch.ones(6, 2) |
| weight = 0.01 |
| bias = 0.2 |
| d_in = 3 |
| d_out = 4 |
| |
| class Foo(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.register_buffer("buf", buf) |
| |
| self.linear = torch.nn.Linear(d_in, d_out) |
| self.linear.weight.data.fill_(weight) |
| self.linear.bias.data.fill_(bias) |
| |
| def forward(self, x): |
| self.buf.add_(5) |
| return self.linear(x).cos() + self.buf.sum() |
| |
| inp = torch.ones(4, 3) |
| ep = torch._export.export( |
| Foo(), |
| (inp,), |
| constraints=[dynamic_dim(inp, 0) >= 3], |
| _functionalize_runtime_assertions=True, |
| ) |
| |
| gs = ep.graph_signature |
| self.assertEqual( |
| gs, |
| ExportGraphSignature( |
| parameters=["L__self___linear.weight", "L__self___linear.bias"], |
| buffers=["L__self___buf"], |
| user_inputs=["arg3_1"], |
| user_outputs=["add_tensor_1"], |
| inputs_to_parameters={ |
| "arg0_1": "L__self___linear.weight", |
| "arg1_1": "L__self___linear.bias", |
| }, |
| inputs_to_buffers={"arg2_1": "L__self___buf"}, |
| buffers_to_mutate={"add_tensor": "L__self___buf"}, |
| backward_signature=None, |
| assertion_dep_token={2: "dep_token7"}, |
| ), |
| ) |
| outputs = next(n for n in ep.graph.nodes if n.op == "output").args[0] |
| self.assertEqual( |
| [str(o) for o in outputs], |
| ["add_tensor", "add_tensor_1", "dep_token7"], |
| ) |
| self.assertEqual( |
| len(outputs), len(gs.buffers_to_mutate) + len(gs.user_outputs) + 1, |
| ) |
| inp = torch.randn(5, 3) |
| self.assertTrue( |
| torch._dynamo.utils.same( |
| # Directly check run output of `ep.graph_module` which is |
| # functionalized. |
| ep.graph_module( |
| torch.full((d_out, d_in), weight), |
| torch.full((d_out,), bias), |
| buf.clone(), |
| inp, |
| ), |
| (buf.add(5), Foo()(inp), torch.empty(0)), |
| ) |
| ) |
| self.assertTrue(torch._dynamo.utils.same(ep(inp), Foo()(inp))) |
| |
| def test_graph_partition_after_assertion_functionalization(self) -> None: |
| def f1(a, b): |
| add = a + b |
| add_1 = add + b |
| add_2 = add_1 + add |
| |
| relu_1 = add_2.relu() # blocked by this |
| |
| add_3 = add_2 + relu_1 |
| add_4 = add_2 + add_3 |
| return add_4, add_2 |
| |
| partitioner1 = CapabilityBasedPartitioner( |
| graph_module=symbolic_trace(f1), |
| operator_support=_AddOperatorSupport(), |
| ) |
| partitions1 = partitioner1.propose_partitions() |
| |
| self.assertEqual( |
| _to_partition_names(partitions1), |
| [{"add_3", "add_4"}, {"add", "add_1", "add_2"}], |
| ) |
| |
| def f2(a, b): |
| add = a + b |
| add_1 = add + b |
| add_2 = add_1 + add |
| |
| assert add_1[0] == 5 |
| |
| relu_1 = add_2.relu() # blocked by this |
| |
| add_3 = add_2 + relu_1 |
| add_4 = add_2 + add_3 |
| return add_4, add_2 |
| |
| inps = (torch.tensor([1, 3, 2]), torch.tensor([2, 3, 4])) |
| gm = export( |
| f2, |
| inps, |
| constraints=[dynamic_dim(inps[0], 0) == dynamic_dim(inps[1], 0)], |
| _functionalize_runtime_assertions=True, |
| ).graph_module |
| partitioner2 = CapabilityBasedPartitioner( |
| graph_module=gm, |
| operator_support=_AtenAddOperatorSupport(), |
| ) |
| partitions2 = partitioner2.propose_partitions() |
| |
| self.assertEqual( |
| _to_partition_names(partitions2), |
| [ |
| {"add_tensor_3", "add_tensor_4"}, |
| {"add_tensor_1", "add_tensor_2", "add_tensor"}, |
| ] |
| ) |
| |
| fused_gm1 = partitioner1.fuse_partitions(partitions1) |
| fused_gm2 = partitioner2.fuse_partitions(partitions2) |
| |
| inps = (torch.tensor([1, 4, 6]), torch.tensor([2, 4, 6])) |
| self.assertTrue( |
| torch._dynamo.utils.same(fused_gm1(*inps)[0], fused_gm2(*inps)[0]), |
| ) |
| |
| # Sub-module `fused_1` is for logic `add = ..., ..., add_2 = ...` |
| output_names1 = _get_output_names(fused_gm1.get_submodule("fused_1")) |
| output_names2 = _get_output_names(fused_gm2.get_submodule("fused_1")) |
| |
| self.assertEqual(output_names1, ["add_2"]) |
| # The extra output `add_tensor_1` is consumed by assertion. |
| self.assertEqual(output_names2, ["add_tensor_1", "add_tensor_2"]) |
| |
| |
| if __name__ == '__main__': |
| run_tests() |
