| # Owner(s): ["module: ProxyTensor"] |
| |
| from torch.testing._internal.common_utils import TestCase, run_tests |
| import torch |
| import torch._dynamo |
| import unittest |
| import warnings |
| import operator |
| from collections.abc import Iterable |
| from torch.nn.utils import stateless |
| from torch.testing._internal.common_device_type import instantiate_device_type_tests |
| from torch.testing._internal.common_methods_invocations import op_db, skip, xfail, skipOps |
| from torch._subclasses.fake_tensor import DynamicOutputShapeException, DataDependentOutputException, FakeTensorMode |
| from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode |
| from torch._decomp import decomposition_table |
| from torch.fx.experimental.symbolic_shapes import ( |
| eval_guards, bind_symbols, fx_placeholder_vals, fx_placeholder_targets, |
| guard_int, GuardOnDataDependentSymNode |
| ) |
| from torch.testing._internal.custom_op_db import custom_op_db |
| from torch.testing._internal.hop_db import hop_db |
| from torch.testing._internal.common_device_type import ops |
| import torch.testing._internal.optests as optests |
| from torch._C import _disabled_torch_function_impl |
| from torch.fx.experimental.proxy_tensor import make_fx, DecompositionInterpreter, get_isolated_graphmodule |
| from torch.utils._pytree import tree_map |
| from torch.fx.passes.runtime_assert import insert_deferred_runtime_asserts |
| from torch import nn |
| import torch._functorch.config |
| import re |
| |
| import functools |
| import itertools |
| |
| aten = torch.ops.aten |
| |
| HAS_CUDA = torch.cuda.is_available() |
| |
| |
| def strip_end(s, suffix): |
| if suffix and s.endswith(suffix): |
| return s[:-len(suffix)] |
| else: |
| return s |
| |
| |
| def show_guards(gm): |
| names = [strip_end(n, "_1") for n in fx_placeholder_targets(gm)] |
| return "\n".join( |
| gm.shape_env.produce_guards(fx_placeholder_vals(gm), names, _simplified=True, input_contexts=None) |
| ) |
| |
| |
| def process_failures(): |
| """ |
| Takes file containing failures like |
| |
| FAILED test/test_proxy_tensor.py::TestProxyTensorOpInfoCPU::test_make_fx_symbolic_exhaustive___getitem___cpu_float32 - RuntimeError: aten.size.default - couldn't find symbolic meta function/decomposition # noqa: B950 |
| |
| and processes them into a list of opinfo xfails |
| """ |
| f = open('pytest_failures') |
| failures = f.readlines() |
| failures = [i.strip() for i in failures] |
| |
| def process_failure_string(s, matcher): |
| out = re.search(matcher, s) |
| return out.groups() |
| |
| SYMBOLIC_TRACE_MATCH = r'exhaustive_(.*)_cpu.*: (.*)' |
| failures = [process_failure_string(s, SYMBOLIC_TRACE_MATCH) for s in failures] |
| |
| def create_normalized_name(op): |
| if op.variant_test_name == '': |
| s = op.name |
| else: |
| s = f"{op.name}.{op.variant_test_name}" |
| return s.replace('.', '_') |
| |
| remap_opinfo = {create_normalized_name(op): (op.name, op.variant_test_name) for op in op_db} |
| |
| print("symbolic_tensor_failures = {") |
| for failure, reason in failures: |
| print(f" xfail{remap_opinfo[failure]}, # {reason}") |
| print("}") |
| |
| |
| USE_TORCHVISION = False |
| try: |
| import torchvision |
| USE_TORCHVISION = True |
| except ImportError: |
| warnings.warn("Couldn't import torchvision. Some of our tests use it, try " |
| "to install it with commands from pytorch.org, post-fixed with " |
| "`--no-deps` to avoid overwriting the pytorch installation", |
| UserWarning) |
| |
| |
| def _create_new_input(x): |
| if not isinstance(x, torch.Tensor): |
| return x |
| if x.dtype != torch.float: |
| return x + 1 |
| if x.is_leaf: |
| return torch.rand_like(x, requires_grad=x.requires_grad) |
| else: |
| return torch.rand_like(x) |
| |
| """ |
| Delays a cos being executed on the unwraptensor until its used. Simulates a CommTensor used |
| """ |
| class UnwrapTensor(torch.Tensor): |
| @staticmethod |
| def __new__(cls, tensor: torch.Tensor): |
| r = torch.Tensor._make_wrapper_subclass( |
| cls, |
| tensor.size(), |
| dtype=tensor.dtype, |
| device=tensor.device, |
| layout=tensor.layout, |
| requires_grad=tensor.requires_grad, |
| ) |
| r._tensor = tensor |
| return r |
| |
| def __repr__(self): |
| # TODO: consider all_gather the local tensors for better debugging |
| return f"UnwrapTensor({self._tensor})" |
| |
| __torch_function__ = _disabled_torch_function_impl |
| |
| @classmethod |
| def __torch_dispatch__(cls, func, types, args=(), kwargs=None): |
| def unwrap(e): |
| ret = e |
| if isinstance(e, UnwrapTensor): |
| ret = e._tensor.cos() |
| |
| return ret |
| |
| args = tree_map(unwrap, args) |
| kwargs = tree_map(unwrap, kwargs) |
| return func(*args, **kwargs) |
| |
| class TestGenericProxyTensor(TestCase): |
| # WARNING: if any of your inputs are index tensors, DO NOT use this |
| # function |
| def _test(self, f, inps): |
| fx_f = make_fx(f, tracing_mode=self.tracing_mode)(*inps) |
| new_inps = tree_map(_create_new_input, inps) |
| r1 = fx_f(*new_inps) |
| r2 = f(*new_inps) |
| self.assertEqual(r1, r2) |
| |
| def test_pre_dispatch_mode_stack(self): |
| def f(a): |
| b = torch.ones(4, 4) |
| return torch.matmul(a, b) |
| # We expect to see matmul in the trace - it should NOT be decomposed into mm. |
| # Also, torch.ones() doesn't show up in the trace. |
| # This is annoying but expected: ones() never dispatches to the Autograd dispatch key, |
| # so our mode never sees it - it goes directly to the BackendSelect key. |
| inp = torch.ones(4, 4) |
| # Test that make_fx(pre_dispatch=True) clears caches properly. |
| from torch._dispatch.python import enable_python_dispatcher |
| with enable_python_dispatcher(): |
| out1 = f(inp) |
| fx_g = make_fx(f, pre_dispatch=True)(inp) |
| self.assertExpectedInline(fx_g.code.strip(), """\ |
| def forward(self, a_1): |
| ones = torch.ops.aten.ones.default([4, 4], device = device(type='cpu'), pin_memory = False) |
| matmul = torch.ops.aten.matmul.default(a_1, ones); a_1 = ones = None |
| return matmul""") |
| |
| def test_pre_dispatch_linear(self): |
| def f(a, b, c): |
| return torch.nn.functional.linear(a, b, c) |
| a = torch.ones(4, 4) |
| b = torch.ones(4, 4) |
| c = torch.ones(4) |
| fx_g = make_fx(f, pre_dispatch=True)(a, b, c) |
| out1 = f(a, b, c) |
| out2 = fx_g(a, b, c) |
| self.assertEqual(out1, out2) |
| |
| def test_pre_dispatch_no_grad(self): |
| def f(a): |
| b = a.sin() |
| torch.set_grad_enabled(False) |
| c = b.cos() |
| torch.set_grad_enabled(True) |
| return b + c.sin() |
| a1 = torch.randn(4, requires_grad=True) |
| a2 = a1.clone().detach().requires_grad_(True) |
| a_tmp = a1.clone().detach().requires_grad_(True) |
| fx_g = make_fx(f, pre_dispatch=True)(a_tmp) |
| out1 = f(a1) |
| out2 = fx_g(a2) |
| self.assertEqual(out1, out2) |
| out1.sum().backward() |
| out2.sum().backward() |
| self.assertEqual(a1.grad, a2.grad) |
| |
| def test_make_fx_simple(self): |
| def f(x): |
| return torch.sin(x) |
| self._test(f, (torch.randn(3),)) |
| |
| def test_scalar_device(self, device='cpu'): |
| def f(a, b): |
| return a + b |
| self._test(f, [torch.randn(3, device=device), torch.tensor(5)]) |
| |
| def test_isolated_graphmodule(self): |
| def is_any_sum(gm): |
| return any(node.target == torch.ops.aten.sum.default for node in gm.graph.nodes) |
| |
| def is_any_digamma(gm): |
| return any(node.target == torch.ops.aten.digamma.default for node in gm.graph.nodes) |
| |
| def is_any_sigmoid(gm): |
| return any(node.target == torch.ops.aten.sigmoid.default for node in gm.graph.nodes) |
| |
| def inner(x): |
| return torch.sum(x) |
| |
| def f(x): |
| gm = get_isolated_graphmodule(inner, (x,), {}) |
| self.assertTrue(is_any_sum(gm)) |
| return x + torch.randn(x.shape) |
| |
| # get_isolated_graphmodule uses make_fx internally that shouldn't be traced |
| # by the outer make_fx call |
| traced = make_fx(f)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| |
| # When factory functions are used, they should not be traced |
| # by the outer make_fx call |
| def inner_with_factory(): |
| val = torch.tensor(float(1)) |
| val.add_(2) |
| return torch.full((10, 10), val).sum() |
| |
| def f1(x): |
| gm = get_isolated_graphmodule(inner_with_factory, (), {}) |
| self.assertTrue(is_any_sum(gm)) |
| return torch.sigmoid(x) |
| |
| def f2(x): |
| gm = get_isolated_graphmodule(f1, (x,), {}) |
| self.assertFalse(is_any_sum(gm)) |
| self.assertTrue(is_any_sigmoid(gm)) |
| return torch.digamma(x) |
| |
| traced = make_fx(f2)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| self.assertFalse(is_any_sigmoid(traced)) |
| self.assertTrue(is_any_digamma(traced)) |
| |
| # Verify nested make_fx calls don't make factory functions to be leaked |
| # into the outer graph. Verify that `make_fx`` itself does not leak its execution. |
| def f2(x): |
| gm = make_fx(f1)(x) |
| self.assertFalse(is_any_sum(gm)) |
| self.assertTrue(is_any_sigmoid(gm)) |
| return torch.digamma(x) |
| |
| traced = make_fx(f2)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| self.assertFalse(is_any_sigmoid(traced)) |
| self.assertTrue(is_any_digamma(traced)) |
| |
| # Verify that the `forward`` function of a graph module produced as a |
| # side effect of an interior `make_fx` is still traced |
| def f3(x): |
| gm = make_fx(f1)(x) |
| self.assertFalse(is_any_sum(gm)) |
| self.assertTrue(is_any_sigmoid(gm)) |
| # `gm.forward`` is still traced |
| return torch.digamma(gm(x)) |
| |
| traced = make_fx(f3)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| self.assertTrue(is_any_sigmoid(traced)) |
| self.assertTrue(is_any_digamma(traced)) |
| |
| # Verify interaction with non-ProxyTensor modes |
| from torch.testing._internal.logging_tensor import LoggingTensorMode |
| |
| def f1_logging(x): |
| with LoggingTensorMode(): |
| gm = get_isolated_graphmodule(inner_with_factory, (), {}) |
| self.assertTrue(is_any_sum(gm)) |
| return torch.sigmoid(x) |
| |
| def f2_logging(x): |
| with LoggingTensorMode(), LoggingTensorMode(): |
| gm = get_isolated_graphmodule(f1_logging, (x,), {}) |
| self.assertFalse(is_any_sum(gm)) |
| self.assertTrue(is_any_sigmoid(gm)) |
| return torch.digamma(x) |
| |
| traced = make_fx(f2_logging)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| self.assertFalse(is_any_sigmoid(traced)) |
| self.assertTrue(is_any_digamma(traced)) |
| |
| # Verify interaction with another tensor subclass |
| # This case currently doesn't work and should raise an error |
| # See: https://github.com/pytorch/pytorch/pull/81764#issuecomment-1200472068 |
| from torch.testing._internal.logging_tensor import LoggingTensor |
| |
| def f1_logging_tensor(x): |
| gm = get_isolated_graphmodule(inner_with_factory, (), {}) |
| self.assertTrue(is_any_sum(gm)) |
| return torch.sigmoid(x) |
| |
| def f2_logging_tensor(x): |
| x = LoggingTensor(x) |
| gm = get_isolated_graphmodule(f1_logging_tensor, (x,), {}) |
| self.assertFalse(is_any_sum(gm)) |
| self.assertTrue(is_any_sigmoid(gm)) |
| return torch.digamma(x) |
| |
| traced = make_fx(f2_logging_tensor)(torch.randn(3)) |
| self.assertFalse(is_any_sum(traced)) |
| self.assertFalse(is_any_sigmoid(traced)) # this fails, sigmoid is traced with LoggingTensor |
| self.assertTrue(is_any_digamma(traced)) |
| |
| # See https://github.com/pytorch/pytorch/issues/97541 |
| def test_empty_like_doesnt_burn_in_defaults(self): |
| def f(x): |
| return torch.empty_like(x) |
| out = make_fx(f)(torch.randn(3)) |
| self.assertExpectedInline(out.code.strip(), """\ |
| def forward(self, x_1): |
| empty_like = torch.ops.aten.empty_like.default(x_1, pin_memory = False); x_1 = None |
| return empty_like""") |
| |
| def test_proxy_tensor_mode_with_decomp_table_preserves_proxy(self): |
| def f(x): |
| y = x.new_zeros(x.size()) |
| y.copy_(x) |
| return y |
| |
| def _new_zeros_decomp(inp, size, dtype=None, layout=None, device=None, pin_memory=None): |
| return torch.zeros(size, dtype=inp.dtype, device=inp.device) |
| |
| factory_func_decomp = {torch.ops.aten.new_zeros.default: _new_zeros_decomp} |
| |
| # When new_zeros() decomposes into torch.zero(), we expect ProxyTensorMode |
| # to still be (re-entrantly) enabled, so that the `torch.zero()` call |
| # returns a ProxyTensor. |
| out = make_fx(f, decomposition_table=factory_func_decomp)(torch.ones(2)) |
| self.assertExpectedInline(out.code, """\ |
| |
| |
| |
| def forward(self, x_1): |
| zeros = torch.ops.aten.zeros.default([2], dtype = torch.float32, device = device(type='cpu'), pin_memory = False) |
| copy_ = torch.ops.aten.copy_.default(zeros, x_1); zeros = x_1 = None |
| return copy_ |
| """) |
| |
| def test_make_fx_reentrant_dispatch(self): |
| def f(x): |
| return torch.ops.aten.norm.Scalar(x, 2.0) |
| |
| def norm_decomp(x, p=2.0): |
| if p != 2.0: |
| raise RuntimeError("can't handle with p != 2") |
| return torch.sqrt(torch.sum(torch.square(x))) |
| |
| decomp = {torch.ops.aten.norm.Scalar: norm_decomp} |
| |
| traced = make_fx(f, decomposition_table=decomp, tracing_mode=self.tracing_mode)(torch.rand(3)) |
| |
| for n in traced.graph.nodes: |
| self.assertTrue("square" not in str(n.target)) |
| self.assertTrue("norm" not in str(n.target)) |
| |
| @unittest.skipIf(not USE_TORCHVISION, "test requires torchvision") |
| def test_resnet18_backward_trace(self): |
| mod = torchvision.models.resnet18() |
| |
| # An old version of this test called the module directly. This works |
| # for tracing_mode == "real", but for fake tensors, we also have to |
| # ensure that the parameters and buffers get wrapped in fake tensors |
| # because free fake tensors are not supported. Fortunately functional_call |
| # does precisely this for us. |
| def f(x, params, buffers): |
| for p in params.values(): |
| p.grad = None |
| loss = torch.func.functional_call(mod, {**params, **buffers}, (x,)).sum() |
| # I could have done this with the functional API, but there is |
| # plenty of exercising this; I want to show mutating API still |
| # works |
| loss.backward() |
| return [p.grad for p in params.values()] |
| |
| inp = torch.randn(3, 3, 250, 250) |
| self._test(f, [inp, dict(mod.named_parameters()), dict(mod.named_buffers())]) |
| |
| def test_varargs(self): |
| def f(*args): |
| return sum(args) |
| |
| self._test(f, [torch.randn(2), torch.randn(2)]) |
| |
| def test_proxy_tensor(self): |
| def f_grad(x): |
| val = x.cos().cos().sum() |
| return torch.autograd.grad(val, x) |
| |
| def f_backward(x): |
| val = x.cos().cos().sum() |
| val.backward() |
| return x.grad |
| |
| for f in [f_grad, f_backward]: |
| self._test(f, [torch.randn(3, requires_grad=True)]) |
| |
| def test_pickle_issue89626(self): |
| import pickle |
| x = torch.randn(2) |
| make_fx(lambda x: x * 2, tracing_mode=self.tracing_mode)(x) |
| pickle.dumps(x) |
| |
| def test_inplace_metadata(self): |
| def f(x): |
| x = x.clone() |
| x.unsqueeze_(-1) |
| assert x.shape[-1] == 1 |
| return x |
| |
| self._test(f, [torch.randn(5)]) |
| |
| def test_mode_tracing_factory_function(self): |
| def f(x): |
| return x + torch.randn(x.shape) |
| |
| # default behavior should trace factory functions |
| traced = make_fx(f, tracing_mode=self.tracing_mode)(torch.randn(3)) |
| self.assertTrue( |
| any( |
| node.target == aten.randn.default |
| for node in traced.graph.nodes |
| ) |
| ) |
| |
| def test_pre_dispatch_functionalization(self): |
| def f(x): |
| a = FunctionalTensorMode(pre_dispatch=True) |
| with a: |
| x_unwrapped = FunctionalTensor.to_functional(x) |
| y = torch.matmul(x_unwrapped, x_unwrapped) |
| y = y + x_unwrapped |
| y.mul_(5) |
| y_unwrapped = torch._from_functional_tensor(y.elem) |
| return y_unwrapped |
| |
| from torch._dispatch.python import enable_python_dispatcher |
| |
| with enable_python_dispatcher(): |
| inp = torch.randn(4, 4) |
| gm = make_fx(f, pre_dispatch=True)(inp) |
| |
| # TODO actually not decompose |
| self.assertExpectedInline(gm.code.strip(), """\ |
| def forward(self, x_1): |
| matmul = torch.ops.aten.matmul.default(x_1, x_1) |
| add = torch.ops.aten.add.Tensor(matmul, x_1); matmul = x_1 = None |
| mul = torch.ops.aten.mul.Tensor(add, 5); add = None |
| return mul""") |
| |
| def test_pre_dispatch_functionalization_view_op(self): |
| def f(x): |
| a = FunctionalTensorMode(pre_dispatch=True) |
| with a: |
| x_unwrapped = FunctionalTensor.to_functional(x) |
| y = torch.matmul(x_unwrapped, x_unwrapped) |
| x_unwrapped = x_unwrapped.transpose(1, 0) |
| y = y + x_unwrapped |
| y = y.view(2, 8) |
| y_unwrapped = torch._from_functional_tensor(y.elem) |
| return y_unwrapped |
| |
| from torch._dispatch.python import enable_python_dispatcher |
| |
| with enable_python_dispatcher(): |
| inp = torch.randn(4, 4) |
| gm = make_fx(f, pre_dispatch=True)(inp) |
| |
| # TODO actually not decompose |
| self.assertExpectedInline(gm.code.strip(), """\ |
| def forward(self, x_1): |
| matmul = torch.ops.aten.matmul.default(x_1, x_1) |
| transpose = torch.ops.aten.transpose.int(x_1, 1, 0); x_1 = None |
| add = torch.ops.aten.add.Tensor(matmul, transpose); matmul = transpose = None |
| view = torch.ops.aten.view.default(add, [2, 8]); add = None |
| return view""") |
| |
| def test_val_metadata_mutation(self): |
| def f(x): |
| y = x.clone() |
| y.unsqueeze_(0) |
| return y |
| |
| traced = make_fx(f, tracing_mode=self.tracing_mode)(torch.randn(3, requires_grad=True)) |
| self.assertEqual([ |
| tuple(node.meta['val'].shape) |
| for node in traced.graph.nodes |
| if 'val' in node.meta |
| ], [(3,), (3,), (1, 3)]) |
| |
| def test_make_fx_overloads(self): |
| def f(x): |
| return x.cos() + torch.randn(x.shape) |
| |
| traced = make_fx(f, tracing_mode=self.tracing_mode)(torch.randn(3)) |
| |
| self.assertTrue(all(isinstance(node.target, torch._ops.OpOverload) |
| for node in traced.graph.nodes if node.op == 'call_function')) |
| |
| def test_tensor_constants(self): |
| def f(): |
| val = torch.tensor(float('inf')) |
| return torch.full((100, 100), val) |
| |
| self._test(f, []) |
| |
| def test_allclose(self): |
| def f(a, b): |
| return torch.allclose(a, b) |
| |
| def test_f(): |
| make_fx(f, tracing_mode=self.tracing_mode)( |
| torch.zeros(3), torch.zeros(3) |
| ) |
| |
| if self.tracing_mode != "real": |
| self.assertRaises(DataDependentOutputException, test_f) |
| else: |
| self.assertRaisesRegex(RuntimeError, "data-dependent", test_f) |
| |
| def test_constant_proxy_tensor_mut(self): |
| def f(): |
| val = torch.tensor(float(1)) |
| val.add_(2) |
| return torch.full((100, 100), val) |
| |
| g = make_fx(f, tracing_mode=self.tracing_mode)() |
| self.assertEqual(g(), f()) |
| # In case we mutated shared state in the g graph! |
| self.assertEqual(g(), f()) |
| |
| def test_constant_unbind(self): |
| def f(): |
| val = torch.tensor([2]) |
| r, = torch.unbind(val, 0) |
| return r.item() |
| |
| g = make_fx(f, tracing_mode=self.tracing_mode)() |
| self.assertEqual(g(), f()) |
| |
| def test_constant_blowup(self): |
| def f(): |
| val = torch.tensor([2]) |
| blowup = val.repeat(1000) |
| return bool(blowup.sum().item() == 2) |
| |
| def test_f(): |
| make_fx(f, tracing_mode=self.tracing_mode)() |
| |
| self.assertRaisesRegex(RuntimeError, "data-dependent", test_f) |
| |
| def test_constant_random(self): |
| def f(): |
| val = torch.tensor([2.0]) |
| val.normal_() |
| return bool(val.item() == 2.1) |
| |
| def test_f(): |
| make_fx(f, tracing_mode=self.tracing_mode)() |
| |
| self.assertRaisesRegex(RuntimeError, "data-dependent", test_f) |
| |
| def test_decomposition_interpreter(self): |
| def fn(x): |
| return torch.nn.functional.silu(x) |
| |
| x = torch.rand((4, 4)) |
| fx_module = make_fx(fn, tracing_mode=self.tracing_mode, decomposition_table=None)(x) |
| |
| found_silu = False |
| for n in fx_module.graph.nodes: |
| if n.target == torch.ops.aten.silu or n.target == torch.ops.aten.silu.default: |
| found_silu = True |
| |
| self.assertTrue(found_silu) |
| |
| new_graph = torch.fx.Graph() |
| silu_decomp_table = {torch.ops.aten.silu.default: decomposition_table[torch.ops.aten.silu.default]} |
| DecompositionInterpreter( |
| fx_module, |
| new_graph=new_graph, |
| decomposition_table=silu_decomp_table, |
| ).run(x) |
| |
| decomposed_module = torch.fx.GraphModule(fx_module, new_graph) |
| |
| for n in decomposed_module.graph.nodes: |
| self.assertTrue(n.target != torch.ops.aten.silu) |
| self.assertTrue(n.target != torch.ops.aten.silu.default) |
| |
| self.assertEqual(fx_module(x), decomposed_module(x)) |
| |
| def test_make_fx_model_fwd_bwd(self): |
| class Foo(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.linear = torch.nn.Linear(5, 5) |
| |
| def forward(self, x): |
| return self.linear(x).relu() |
| |
| model = Foo() |
| |
| def f(x, params): |
| out = torch.func.functional_call(model, params, x).sum() |
| out.backward() |
| return list(params.values()) |
| input = torch.randn(3, 5, requires_grad=True) |
| params = dict(model.named_parameters()) |
| fx_f = make_fx(f, tracing_mode=self.tracing_mode)(input, params) |
| # fx may change the order of parameters in list, so using set() to compare |
| self.assertTrue( |
| torch.allclose(fx_f(input, params)[0], f(input, params)[0]) |
| or |
| torch.allclose(fx_f(input, params)[0], f(input, params)[1]) |
| ) |
| self.assertTrue( |
| torch.allclose(fx_f(input, params)[1], f(input, params)[0]) |
| or |
| torch.allclose(fx_f(input, params)[1], f(input, params)[1]) |
| ) |
| |
| def test_make_fx_model_double_param(self): |
| class Emformer(torch.nn.Module): |
| def __init__( |
| self, |
| input_dim: int = 256, |
| ) -> None: |
| super().__init__() |
| |
| self.layer_norm = torch.nn.LayerNorm(input_dim) |
| |
| def forward(mod_self, x): # noqa: B902 |
| self.assertTrue(isinstance(mod_self.layer_norm.weight, torch.Tensor)) |
| y = mod_self.layer_norm(x) |
| self.assertTrue(isinstance(mod_self.layer_norm.weight, torch.Tensor)) |
| z = mod_self.layer_norm(y) |
| return z |
| |
| |
| gm = make_fx(Emformer())(torch.randn(16, 1, 256)) |
| ops = {n.target for n in gm.graph.nodes if n.op == 'call_function'} |
| self.assertEqual(len(ops), 2) |
| |
| |
| def test_make_fx_model_fwd_bwd_wgtupdate(self): |
| class Foo(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.linear = torch.nn.Linear(5, 5) |
| |
| def forward(self, x): |
| return self.linear(x).relu() |
| |
| model = Foo() |
| |
| def f(args, params, buffers): |
| for p in params.values(): |
| p.grad = None |
| if not isinstance(args, Iterable): |
| args = [args] |
| params_and_buffers = {**params, **buffers} |
| out = torch.func.functional_call(model, params_and_buffers, args) |
| out.sum().backward() |
| return [p - 1e-4 * p.grad for p in params.values()] |
| |
| input = torch.randn(3, 5, requires_grad=True) |
| params = dict(model.named_parameters()) |
| buffers = dict(model.named_buffers()) |
| fx_f = make_fx(f, tracing_mode=self.tracing_mode)(input, params, buffers) |
| # fx may change the order of parameters in list, so using set() to compare |
| # also there is a numerical difference in results so changing atol from 1e-08 to 1e-03 |
| self.assertTrue( |
| torch.allclose(fx_f(input, params, buffers)[0], f(input, params, buffers)[0], atol=1e-03) |
| or |
| torch.allclose(fx_f(input, params, buffers)[0], f(input, params, buffers)[1], atol=1e-03) |
| ) |
| self.assertTrue( |
| torch.allclose(fx_f(input, params, buffers)[1], f(input, params, buffers)[0], atol=1e-03) |
| or |
| torch.allclose(fx_f(input, params, buffers)[1], f(input, params, buffers)[1], atol=1e-03) |
| ) |
| |
| def test_trace_subclasses(self): |
| def f1(x): |
| x = UnwrapTensor(x) |
| y = x * 2 |
| return y |
| |
| def f2(x): |
| wrapped = UnwrapTensor(x) |
| y = x * wrapped |
| return y |
| |
| inp = [torch.randn(5)] |
| self._test(f1, inp) |
| self._test(f2, inp) |
| |
| def test_partial_decomp(self): |
| def f(a, b, c): |
| x = torch.addmm(a, b, c) |
| y = torch.addmm(a, b, c, beta=2, alpha=1) |
| return x + y |
| inps = [torch.randn(5, 5), torch.randn(5, 5), torch.randn(5, 5)] |
| fx_g = make_fx(f)(*inps) |
| |
| def addmm(a, b, c, beta=1, alpha=1): |
| if beta == 1 and alpha == 1: |
| return NotImplemented |
| return beta * a + alpha * (b @ c) |
| |
| decomposed_fx = make_fx(f, decomposition_table={aten.addmm.default: addmm})(*inps) |
| |
| self.assertEqual(fx_g(*inps), decomposed_fx(*inps)) |
| self.assertEqual(len([n for n in fx_g.graph.nodes if n.target == aten.addmm.default]), 2) |
| self.assertEqual(len([n for n in decomposed_fx.graph.nodes if n.target == aten.addmm.default]), 1) |
| |
| def test_decomp_of_capture(self): |
| val = torch.randn(5) |
| |
| def f(x): |
| return x.t() + val.t() |
| |
| def nop(x): |
| return x.cos() |
| |
| traced = make_fx(f, decomposition_table={torch.ops.aten.t.default: nop})(torch.randn(5)) |
| self.assertEqual(len([n for n in traced.graph.nodes if n.target == torch.ops.aten.t.default]), 0) |
| |
| |
| @unittest.skipIf(not HAS_CUDA, 'CUDA-only test') |
| def test_amp_cache(self): |
| layer = torch.nn.Conv2d(3, 3, 3).cuda() |
| |
| def f(x, w): |
| return torch.nn.functional.conv2d(x, w, stride=layer.stride) |
| |
| inp = torch.randn(4, 3, 10, 10, device='cuda') |
| with torch.autocast('cuda'): |
| out_graph = make_fx(f)(inp, layer.weight).graph |
| out_graph2 = make_fx(f)(inp, layer.weight).graph |
| |
| self.assertEqual(len(out_graph.nodes), len(out_graph2.nodes)) |
| for a, b in zip(out_graph.nodes, out_graph2.nodes): |
| self.assertEqual(a.op, b.op) |
| |
| def test_strides(self): |
| def f(x): |
| self.assertTrue(x.is_contiguous()) |
| self.assertFalse(x.is_contiguous(memory_format=torch.channels_last)) |
| x = x.permute(0, 3, 1, 2) |
| self.assertFalse(x.is_contiguous()) |
| self.assertTrue(x.is_contiguous(memory_format=torch.channels_last)) |
| return x |
| make_fx(f)(torch.randn(2, 3, 4, 5)) |
| |
| def f(x): |
| self.assertTrue(x.is_contiguous()) |
| y = x[:, 1] |
| self.assertFalse(y.is_contiguous()) |
| y = x[:, ::2] |
| self.assertFalse(y.is_contiguous()) |
| return x.cos() |
| |
| make_fx(f)(torch.randn(2, 3, 4, 5)) |
| |
| def test_pr_86917(self): |
| # Tests the issue brought up here https://github.com/pytorch/pytorch/pull/86917#issuecomment-1283155344 |
| def f(a, b): |
| return torch.ops.aten.nll_loss_forward(a, b, None, 1, 10) |
| |
| self._test(f, [torch.randn(1, 10), torch.zeros(1, dtype=torch.long)]) |
| |
| class TestGenericProxyTensorReal(TestGenericProxyTensor): |
| tracing_mode = "real" |
| |
| |
| class TestGenericProxyTensorFake(TestGenericProxyTensor): |
| tracing_mode = "fake" |
| |
| |
| class TestGenericProxyTensorSymbolic(TestGenericProxyTensor): |
| tracing_mode = "symbolic" |
| |
| |
| del TestGenericProxyTensor |
| |
| |
| class TestRealProxyTensor(TestCase): |
| def test_error_on_data_dependent_ops(self): |
| def f(): |
| x = torch.randn([]) |
| y = torch.randn([]) |
| assert torch.allclose(x * y, y * x) |
| z = float(x) |
| z2 = float(y) |
| |
| # Smoke tests |
| make_fx(f, _error_on_data_dependent_ops=False)() |
| make_fx(f, pre_dispatch=True, _error_on_data_dependent_ops=False)() |
| |
| class TestFakeProxyTensor(TestCase): |
| def test_issue82547(self): |
| x = nn.Parameter(torch.randn(3, 3)) |
| |
| def f(): |
| return torch.ops.aten.t.default(x) |
| self.assertRaisesRegex(Exception, "Please convert all Tensors", lambda: make_fx(f, tracing_mode="fake")()) |
| |
| class A(torch.Tensor): |
| pass |
| |
| x = A(torch.randn(3, 3)) |
| self.assertRaisesRegex(TypeError, "Multiple dispatch failed", lambda: make_fx(f, tracing_mode="fake")()) |
| |
| def test_use_fake_and_tensor(self): |
| def f(x, y): |
| z = torch.tensor([2.0, 3.0]) |
| return x + y + z |
| |
| g = make_fx(f, tracing_mode="fake")(torch.randn(2), torch.randn(2)) |
| x, y = torch.randn(2), torch.randn(2) |
| self.assertEqual(g(x, y), f(x, y)) |
| |
| def test_free_fake(self): |
| def f(x): |
| return torch.add(x, y) |
| |
| with FakeTensorMode() as fake_mode: |
| y = torch.randn(2) |
| make_fx(f, tracing_mode="real")(torch.randn(2)) |
| |
| def test_fused_adam(self): |
| # See https://github.com/pytorch/pytorch/issues/99356 |
| params = [torch.randn(10, 10) for _ in range(10)] |
| grads = [torch.randn(10, 10) for _ in range(10)] |
| exp_avgs = [torch.randn(10, 10) for _ in range(10)] |
| exp_avg_sqs = [torch.randn(10, 10) for _ in range(10)] |
| max_exp_avg_sqs = [torch.randn(10, 10) for _ in range(10)] |
| state_steps = [torch.tensor(0) for _ in range(10)] |
| |
| def fused_adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps): |
| (new_params, _, _, _, _) = aten._fused_adam.default( |
| params, |
| grads, |
| exp_avgs, |
| exp_avg_sqs, |
| max_exp_avg_sqs, |
| state_steps, |
| lr=0.1, |
| beta1=0.9, |
| beta2=0.999, |
| weight_decay=0.01, |
| eps=1e-8, |
| amsgrad=False, |
| maximize=False, |
| ) |
| |
| for p, new_p in zip(params, new_params): |
| p.copy_(new_p) |
| |
| return params |
| |
| gm = make_fx(fused_adam, tracing_mode='fake')( |
| params, |
| grads, |
| exp_avgs, |
| exp_avg_sqs, |
| max_exp_avg_sqs, |
| state_steps, |
| ) |
| ensure_ops_have_val = [aten._fused_adam.default, operator.getitem] |
| for n in gm.graph.nodes: |
| if n.op == "call_function" and n.target in ensure_ops_have_val: |
| self.assertIn('val', n.meta) |
| |
| def test_alias(self): |
| def f(x): |
| return torch.ops.aten.alias(x) |
| |
| r = str(make_fx(f, tracing_mode="fake")(torch.randn(2)).code).strip() |
| # NB: this should not have a detach call |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1): |
| alias = torch.ops.aten.alias.default(x_1); x_1 = None |
| return alias""") |
| |
| def test_meta(self): |
| def f(x): |
| a = x.cos() |
| b = torch.var_mean(a, dim=0) |
| c = b * 2 |
| return c |
| |
| out = make_fx(f, tracing_mode="fake")(torch.randn(5, 5)) |
| for n in out.graph.nodes: |
| if n.op == 'output': |
| continue |
| self.assertTrue('val' in n.meta) |
| |
| def _get_node(fx_g, cond): |
| for n in fx_g.graph.nodes: |
| if cond(n): |
| return n |
| raise AssertionError |
| |
| def _get_free_symbols(shape_env): |
| vars = tuple(shape_env.var_to_val.keys()) |
| return len([var for var in vars if var not in shape_env.replacements]) |
| |
| def _trace(f, *args): |
| inps = [torch.randn(arg) for arg in args] |
| return make_fx(f, tracing_mode="symbolic")(*inps) |
| |
| # TODO: Need to test the guards themselves specifically as well |
| class TestSymbolicTracing(TestCase): |
| def _test_dynamic(self, fn, trace_inputs, test_inputs, assert_eq=True): |
| """ |
| Tests fn traced with trace_inputs against test_inputs |
| Also returns shape env |
| """ |
| trace_inputs = [torch.randn(shape) for shape in trace_inputs] |
| traced_f = make_fx(fn, tracing_mode="symbolic")(*trace_inputs) |
| for input in test_inputs: |
| input = [torch.randn(shape) for shape in input] |
| rx, ry = traced_f(*input), fn(*input) |
| if assert_eq: |
| self.assertEqual(rx, ry) |
| return traced_f |
| |
| |
| def test_debug_interpreter(self): |
| import torch.library |
| from torch.library import Library |
| |
| foo = Library("foo", "DEF") # noqa: TOR901 |
| foo.define("foo(Tensor self) -> Tensor") |
| |
| # Operator where meta and cpu disagree on strides |
| @torch.library.impl(foo, "foo", "CPU") |
| def foo_cpu(x): |
| return x.clone().T |
| |
| @torch.library.impl(foo, "foo", "Meta") |
| def foo_meta(x): |
| return x.clone() |
| |
| def f(x): |
| return torch.ops.foo.foo.default(x) |
| |
| gm = make_fx(f, tracing_mode="symbolic")(torch.randn(2, 2)) |
| from torch._functorch.compilers import DebugInterpreter |
| |
| interp = DebugInterpreter(gm) |
| |
| # input mismatch is caught (indicates guard problem) |
| self.assertRaisesRegex( |
| AssertionError, r"3 != 1", |
| lambda: interp.run(torch.randn(3, 3).T), |
| ) |
| |
| # Catch the incorrect meta |
| self.assertRaisesRegex( |
| AssertionError, r"\(3, 1\) != \(1, 3\)", |
| lambda: interp.run(torch.randn(3, 3)) |
| ) |
| |
| def test_int_input(self): |
| def f(x, y): |
| return x.view(y) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(3, 4), 12).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| view = torch.ops.aten.view.default(x_1, [y_1]); x_1 = y_1 = None |
| return view""") |
| |
| def test_resize_from_zero(self): |
| def f(x, y): |
| x.resize_(y.size(0)) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(0), torch.empty(2)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| sym_size_int = torch.ops.aten.sym_size.int(y_1, 0); y_1 = None |
| resize_ = torch.ops.aten.resize_.default(x_1, [sym_size_int]); x_1 = sym_size_int = resize_ = None |
| return None""") |
| |
| def test_broadcast_shapes(self): |
| def f(x, y): |
| return torch.functional.broadcast_shapes(x.size(), y.size()[0]) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(3, 1), torch.empty(5)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| sym_size_int = torch.ops.aten.sym_size.int(x_1, 0); x_1 = None |
| sym_size_int_1 = torch.ops.aten.sym_size.int(y_1, 0); y_1 = None |
| return (sym_size_int, sym_size_int_1)""") |
| |
| def test_deduped_shape(self): |
| def f(s0, s1, x, y): |
| return torch.functional.broadcast_shapes(x.size(), y.size()[0]), torch.empty(x.shape[0]) |
| |
| x = torch.empty(3, 1) |
| y = torch.empty(5) |
| from torch.fx.experimental.symbolic_shapes import ShapeEnv |
| shape_env = ShapeEnv() |
| |
| with FakeTensorMode(shape_env=shape_env, static_shapes=False) as fake_mode: |
| x = fake_mode.from_tensor(x) |
| y = fake_mode.from_tensor(y) |
| r = str(make_fx(f, tracing_mode="real")(x.shape[0], y.shape[0], x, y).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, s0_1, s1_1, x_1, y_1): |
| empty = torch.ops.aten.empty.memory_format([s0_1], device = device(type='cpu'), pin_memory = False) |
| return ((s0_1, s1_1), empty)""") |
| |
| def test_non_deduped_shape(self): |
| def f(x, y): |
| return torch.functional.broadcast_shapes(x.size(), y.size()[0]), torch.empty(x.shape[0]) |
| |
| x = torch.empty(3, 1) |
| y = torch.empty(5) |
| from torch.fx.experimental.symbolic_shapes import ShapeEnv |
| shape_env = ShapeEnv() |
| |
| with FakeTensorMode(shape_env=shape_env, static_shapes=False) as fake_mode: |
| x = fake_mode.from_tensor(x) |
| y = fake_mode.from_tensor(y) |
| r = str(make_fx(f, tracing_mode="real")(x, y).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| sym_size_int = torch.ops.aten.sym_size.int(x_1, 0); x_1 = None |
| sym_size_int_1 = torch.ops.aten.sym_size.int(y_1, 0); y_1 = None |
| empty = torch.ops.aten.empty.memory_format([sym_size_int], device = device(type='cpu'), pin_memory = False) |
| return ((sym_size_int, sym_size_int_1), empty)""") |
| |
| def test_unary(self): |
| def f(x): |
| assert x.shape[0] < 20 |
| return x.cos() |
| test_inputs = [] |
| test_inputs.append([(2, 5)]) |
| test_inputs.append([(6, 8)]) |
| gm = self._test_dynamic(f, [(3, 4)], test_inputs) |
| self.assertTrue(eval_guards(gm, torch.randn(4, 5))) |
| self.assertEqual(repr(bind_symbols(gm, torch.randn(4, 5))), "{s0: 4, s1: 5}") |
| self.assertFalse(eval_guards(gm, torch.randn(25, 5))) |
| self.assertExpectedInline(show_guards(gm), """L['x'].size()[0] <= 19""") |
| |
| def test_repeat_interleave(self): |
| def f(src_tokens, beam_size_src): |
| return src_tokens.repeat_interleave(beam_size_src.size(0), 0) |
| |
| prompt_size = 64 |
| vocab_size = 64 |
| batch_size = 4 |
| src_tokens = torch.randint(1, vocab_size, (batch_size, prompt_size)) |
| gm = make_fx(f, tracing_mode="symbolic")(src_tokens, torch.randn(5)) |
| self.assertEqual(len(gm.shape_env.guards), 0) |
| |
| def test_non_symint_size_spec(self): |
| # this isn't really a proxy tensor test, but it's the most convenient |
| # way to get a fake tensor with symbolic sizes |
| def f(x): |
| torch._C._non_sym_sizes(x) |
| return x + 1 |
| |
| x = torch.randn(2, 3) |
| make_fx(f, tracing_mode="symbolic")(x) |
| |
| # https://github.com/pytorch/pytorch/issues/108195 |
| def test_symbolic_repeat_interleave(self): |
| def f(y, x): |
| return y.repeat_interleave(x, dim=1) |
| |
| y = torch.tensor([[1, 2], [3, 4]]) |
| x = torch.tensor([2, 3]) |
| r = str(make_fx(f, tracing_mode="symbolic")(y, x).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, y_1, x_1): |
| repeat_interleave = torch.ops.aten.repeat_interleave.Tensor(x_1); x_1 = None |
| index_select = torch.ops.aten.index_select.default(y_1, 1, repeat_interleave); y_1 = repeat_interleave = None |
| return index_select""") |
| |
| def test_mod_gcd_unbacked(self): |
| def f(_a, _b, _stride): |
| a = _a.item() |
| b = _b.item() |
| stride = _stride.item() |
| torch._check_is_size(a) |
| torch._check_is_size(b) |
| torch._check_is_size(stride) |
| ta = torch.randn(a * stride) |
| tb = torch.randn(b * stride) |
| r = torch.cat([ta, tb]) |
| return r.view(a + b, stride) |
| |
| _a = torch.tensor(30) |
| _b = torch.tensor(20) |
| _stride = torch.tensor(10) |
| r = str(make_fx(f, tracing_mode="symbolic")(_a, _b, _stride).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, _a_1, _b_1, _stride_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(_a_1); _a_1 = None |
| _local_scalar_dense_1 = torch.ops.aten._local_scalar_dense.default(_b_1); _b_1 = None |
| _local_scalar_dense_2 = torch.ops.aten._local_scalar_dense.default(_stride_1); _stride_1 = None |
| mul = _local_scalar_dense * _local_scalar_dense_2 |
| randn = torch.ops.aten.randn.default([mul], device = device(type='cpu'), pin_memory = False); mul = None |
| mul_1 = _local_scalar_dense_1 * _local_scalar_dense_2 |
| randn_1 = torch.ops.aten.randn.default([mul_1], device = device(type='cpu'), pin_memory = False); mul_1 = None |
| cat = torch.ops.aten.cat.default([randn, randn_1]); randn = randn_1 = None |
| add = _local_scalar_dense + _local_scalar_dense_1; _local_scalar_dense = _local_scalar_dense_1 = None |
| view = torch.ops.aten.view.default(cat, [add, _local_scalar_dense_2]); cat = add = _local_scalar_dense_2 = None |
| return view""") |
| |
| def test_cumsum_unbacked(self): |
| def f(x): |
| y = x.item() |
| z = torch.randn((3, y, 3)) |
| return z.cumsum(0) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.tensor([5])).code).strip() |
| self.assertExpectedInline( |
| r, """\ |
| def forward(self, x_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(x_1); x_1 = None |
| randn = torch.ops.aten.randn.default([3, _local_scalar_dense, 3], device = device(type='cpu'), pin_memory = False); _local_scalar_dense = None |
| cumsum = torch.ops.aten.cumsum.default(randn, 0); randn = None |
| return cumsum""" # noqa: B950 |
| ) |
| |
| |
| def test_repeat_interleave_unbacked_output_size(self): |
| def f(x, y): |
| s = x.sum().item() |
| return y.repeat_interleave(x, dim=0, output_size=s) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.tensor([2, 3]), torch.randn(2)).code).strip() |
| self.assertExpectedInline( |
| r, """\ |
| def forward(self, x_1, y_1): |
| sum_1 = torch.ops.aten.sum.default(x_1) |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(sum_1); sum_1 = None |
| repeat_interleave = torch.ops.aten.repeat_interleave.Tensor(x_1, output_size = _local_scalar_dense); x_1 = _local_scalar_dense = None |
| index_select = torch.ops.aten.index_select.default(y_1, 0, repeat_interleave); y_1 = repeat_interleave = None |
| return index_select""" # noqa: B950 |
| ) |
| |
| def test_arange_unbacked_output_size(self): |
| def f(x): |
| return torch.arange(0, x) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.tensor(10)).code).strip() |
| self.assertExpectedInline( |
| r, """\ |
| def forward(self, x_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(x_1); x_1 = None |
| arange = torch.ops.aten.arange.start(0, _local_scalar_dense, device = device(type='cpu'), pin_memory = False); _local_scalar_dense = None |
| return arange""" # noqa: B950 |
| ) |
| |
| def test_adv_index_batch(self): |
| def f(src_tokens): |
| bsz, src_len = src_tokens.size()[:2] |
| start_step = src_tokens.shape[1] |
| beam_size = 1 |
| generate_size = 64 |
| max_len = src_len + generate_size |
| tokens = torch.zeros(bsz * beam_size, max_len).to(src_tokens).long().fill_(0) |
| tokens[:, :start_step] = src_tokens.repeat_interleave(beam_size, 0) |
| return tokens |
| |
| prompt_size = 64 |
| vocab_size = 64 |
| batch_size = 4 |
| src_tokens = torch.randint(1, vocab_size, (batch_size, prompt_size)) |
| gm = make_fx(f, tracing_mode="symbolic")(src_tokens) |
| # Guards to rule out batch_size == sys.maxsize (wobbling between 2 and |
| # 1 ok) |
| self.assertEqual(len(gm.shape_env.guards), 1) |
| |
| @unittest.skipIf(not HAS_CUDA, 'CUDA-only test') |
| def test_cpu_scalar_cuda(self): |
| # Extracted from wave2vec2 |
| def f(a, b): |
| return (a * b) @ b |
| |
| r = str( |
| make_fx(f, tracing_mode="symbolic")( |
| torch.tensor(1.0), torch.randn(2, 2, device='cuda') |
| ).code |
| ).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1, b_1): |
| mul = torch.ops.aten.mul.Tensor(a_1, b_1); a_1 = None |
| mm = torch.ops.aten.mm.default(mul, b_1); mul = b_1 = None |
| return mm""") |
| |
| def test_binary_broadcast(self): |
| def f(a, b): |
| c = a * b |
| return c |
| |
| test_inputs = [] |
| test_inputs.append([(1, 5), (3, 1)]) |
| test_inputs.append([(1, 4), (4, 1)]) |
| shape_env = self._test_dynamic(f, [(1, 2), (3, 1)], test_inputs).shape_env |
| assert len(shape_env.guards) == 0 |
| |
| def test_multiply_shape(self): |
| def f(a): |
| return torch.empty(a.shape[0] * 2) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(4)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| sym_size_int = torch.ops.aten.sym_size.int(a_1, 0); a_1 = None |
| mul = sym_size_int * 2; sym_size_int = None |
| empty = torch.ops.aten.empty.memory_format([mul], device = device(type='cpu'), pin_memory = False); mul = None |
| return empty""") |
| |
| def test_item(self): |
| def f(a): |
| r = a.item() |
| return r * a |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.randn(1)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(a_1) |
| mul = torch.ops.aten.mul.Tensor(a_1, _local_scalar_dense); a_1 = _local_scalar_dense = None |
| return mul""") |
| |
| def test_tensor_symfloat(self): |
| def f(a): |
| r = torch.tensor(a.size(0) ** 2.0) |
| assert r.dtype is torch.float |
| return r |
| |
| gm = make_fx(f, tracing_mode="symbolic")(torch.randn(2)) |
| r = str(gm.code).strip() |
| # NB: this specializes, which is fine, the point is to make sure the |
| # dtype inference is correct |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| _tensor_constant0 = self._tensor_constant0 |
| lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None |
| return lift_fresh_copy""") |
| self.assertEqual(gm._tensor_constant0, torch.tensor(4.0)) |
| |
| def test_item_to_constructor(self): |
| def f(a): |
| r = a.item() |
| return torch.empty(r) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.randint(5, (1,))).code).strip() |
| self.assertExpectedInline( |
| r, """\ |
| def forward(self, a_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(a_1); a_1 = None |
| empty = torch.ops.aten.empty.memory_format([_local_scalar_dense], device = device(type='cpu'), pin_memory = False); _local_scalar_dense = None |
| return empty""" # noqa: B950 |
| ) |
| |
| |
| def test_setitem_symint(self): |
| # from moco |
| # https://github.com/pytorch/pytorch/issues/101939 |
| def f(x): |
| x[0] = x.size(0) |
| return x |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.randn(10)).code).strip() |
| self.assertExpectedInline( |
| r, """\ |
| def forward(self, x_1): |
| sym_size_int = torch.ops.aten.sym_size.int(x_1, 0) |
| scalar_tensor = torch.ops.aten.scalar_tensor.default(sym_size_int, dtype = torch.float32, layout = torch.strided, device = device(type='cpu')); sym_size_int = None |
| select = torch.ops.aten.select.int(x_1, 0, 0) |
| copy_ = torch.ops.aten.copy_.default(select, scalar_tensor); select = scalar_tensor = copy_ = None |
| return x_1""" # noqa: B950 |
| ) |
| |
| def test_dynamic_pointwise_scalar(self): |
| def f(gravity, mask): |
| gravity[mask, 0] = gravity[mask, 0] * -1 |
| |
| r = str(make_fx(f, tracing_mode="symbolic")( |
| torch.randn((12, 4)), |
| torch.randint(0, 2, (12,), dtype=torch.bool) |
| ).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, gravity_1, mask_1): |
| select = torch.ops.aten.select.int(gravity_1, 1, 0) |
| index = torch.ops.aten.index.Tensor(select, [mask_1]); select = None |
| mul = torch.ops.aten.mul.Tensor(index, -1); index = None |
| select_1 = torch.ops.aten.select.int(gravity_1, 1, 0); gravity_1 = None |
| index_put_ = torch.ops.aten.index_put_.default(select_1, [mask_1], mul); select_1 = mask_1 = mul = index_put_ = None |
| return None""") |
| |
| def test_reflect_r_over_x(self): |
| def reflect_R_over_x(R): |
| reflect = torch.eye(3, device=R.device) |
| reflect[0, 0] = -1 |
| return reflect @ R @ reflect |
| |
| def f(crop_camera, mask): |
| crop_camera[mask] = reflect_R_over_x(crop_camera[mask]) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")( |
| torch.randn((12, 3, 3)), |
| torch.randint(0, 2, (12,), dtype=torch.bool) |
| ).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, crop_camera_1, mask_1): |
| index = torch.ops.aten.index.Tensor(crop_camera_1, [mask_1]) |
| eye = torch.ops.aten.eye.default(3, device = device(type='cpu'), pin_memory = False) |
| _tensor_constant0 = self._tensor_constant0 |
| lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None |
| select = torch.ops.aten.select.int(eye, 0, 0) |
| select_1 = torch.ops.aten.select.int(select, 0, 0); select = None |
| copy_ = torch.ops.aten.copy_.default(select_1, lift_fresh_copy); select_1 = lift_fresh_copy = copy_ = None |
| sym_size_int = torch.ops.aten.sym_size.int(index, 0) |
| expand = torch.ops.aten.expand.default(eye, [sym_size_int, 3, 3]) |
| view = torch.ops.aten.view.default(expand, [sym_size_int, 3, 3]); expand = None |
| sym_size_int_1 = torch.ops.aten.sym_size.int(crop_camera_1, 1) |
| sym_size_int_2 = torch.ops.aten.sym_size.int(crop_camera_1, 2) |
| expand_1 = torch.ops.aten.expand.default(index, [sym_size_int, sym_size_int_1, sym_size_int_2]); index = None |
| view_1 = torch.ops.aten.view.default(expand_1, [sym_size_int, sym_size_int_1, sym_size_int_2]); expand_1 = sym_size_int_1 = sym_size_int_2 = None |
| bmm = torch.ops.aten.bmm.default(view, view_1); view = view_1 = None |
| view_2 = torch.ops.aten.view.default(bmm, [sym_size_int, 3, 3]); bmm = None |
| mul_4 = sym_size_int * 3 |
| view_3 = torch.ops.aten.view.default(view_2, [mul_4, 3]); view_2 = mul_4 = None |
| mm = torch.ops.aten.mm.default(view_3, eye); view_3 = eye = None |
| view_4 = torch.ops.aten.view.default(mm, [sym_size_int, 3, 3]); mm = sym_size_int = None |
| index_put_ = torch.ops.aten.index_put_.default(crop_camera_1, [mask_1], view_4); crop_camera_1 = mask_1 = view_4 = index_put_ = None |
| return None""") # noqa: B950 |
| |
| def test_unbacked_slice(self): |
| def f(x, m): |
| x = x[m] |
| return x[slice(None, None, None), slice(None, None, None), slice(None, 2, None)] |
| |
| make_fx(f, tracing_mode="symbolic")( |
| torch.randn((12, 3, 3)), |
| torch.randint(0, 2, (12,), dtype=torch.bool) |
| ) |
| |
| @unittest.skipIf(not USE_TORCHVISION, "test requires torchvision") |
| def test_unbacked_batch_resnet(self): |
| mod = torchvision.models.resnet18() |
| |
| def f(x, mask, params, buffers): |
| for p in itertools.chain([x, mask], params.values(), buffers.values()): |
| for s in p.shape: |
| guard_int(s) |
| x = x[mask] |
| torch._check(x.shape[0] >= 1) |
| for p in params.values(): |
| p.grad = None |
| return torch.func.functional_call(mod, {**params, **buffers}, (x,)).sum() |
| |
| make_fx(f, tracing_mode="symbolic")( |
| torch.randn(3, 3, 250, 250), |
| torch.randint(0, 2, (3,), dtype=torch.bool), |
| dict(mod.named_parameters()), |
| dict(mod.named_buffers()), |
| ) |
| |
| def test_boolean_index(self): |
| def f(images, handedness, valid): |
| images = images[valid] |
| handedness = handedness[valid] |
| right_hand_mask = handedness == 1 |
| images[right_hand_mask] = images[right_hand_mask].flip(-1) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")( |
| torch.randint(0, 256, (512, 1, 96, 96)), |
| torch.randint(0, 1, (512,)), |
| torch.randint(0, 2, (512,), dtype=torch.bool) |
| ).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, images_1, handedness_1, valid_1): |
| index = torch.ops.aten.index.Tensor(images_1, [valid_1]); images_1 = None |
| index_1 = torch.ops.aten.index.Tensor(handedness_1, [valid_1]); handedness_1 = valid_1 = None |
| eq = torch.ops.aten.eq.Scalar(index_1, 1); index_1 = None |
| index_2 = torch.ops.aten.index.Tensor(index, [eq]) |
| flip = torch.ops.aten.flip.default(index_2, [-1]); index_2 = None |
| index_put_ = torch.ops.aten.index_put_.default(index, [eq], flip); index = eq = flip = index_put_ = None |
| return None""") |
| |
| def test_neg_shape(self): |
| def f(a): |
| return torch.empty(-a.shape[0] + 10) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(2)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| sym_size_int = torch.ops.aten.sym_size.int(a_1, 0); a_1 = None |
| neg = -sym_size_int; sym_size_int = None |
| add = neg + 10; neg = None |
| empty = torch.ops.aten.empty.memory_format([add], device = device(type='cpu'), pin_memory = False); add = None |
| return empty""") |
| |
| def test_unbacked_unification(self): |
| def f(x, y): |
| z = torch.zeros(x.item()) |
| return z + y |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.tensor(10), torch.randn(10)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(x_1); x_1 = None |
| zeros = torch.ops.aten.zeros.default([_local_scalar_dense], device = device(type='cpu'), pin_memory = False); _local_scalar_dense = None |
| add = torch.ops.aten.add.Tensor(zeros, y_1); zeros = y_1 = None |
| return add""") # noqa: B950 |
| |
| def test_reshape_divisibility_unbacked(self): |
| def f(x): |
| i0 = x.item() |
| r = torch.zeros(i0, 4, 20) |
| r = r.transpose(2, 1) |
| return r.reshape(-1, 80) |
| make_fx(f, tracing_mode="symbolic")(torch.tensor(24)) |
| |
| def test_view_divisibility_unbacked(self): |
| def f(x): |
| i0 = x.item() |
| r = torch.zeros(i0, 192) |
| return r.view(12, -1, 192) |
| make_fx(f, tracing_mode="symbolic")(torch.tensor(24)) |
| |
| @unittest.skipIf(not HAS_CUDA, 'CUDA-only test') |
| def test_view_divisibility_unbacked_relatively_prime(self): |
| # See https://github.com/pytorch/pytorch/issues/123651 |
| def f(x): |
| i0 = x.item() |
| torch._check_is_size(i0) |
| # To trigger the original issue, the max bound has to |
| # be chosen such that 448 / 447 < 2 (which it is.) |
| torch._check(i0 <= 448) |
| return torch.zeros(256 * i0).view(-1, 447) |
| make_fx(f, tracing_mode="symbolic")(torch.tensor(256 * 447, device="cuda")) |
| |
| def test_unbacked_unify_guard(self): |
| def f(x, y): |
| z = torch.zeros(x.item()) |
| torch._check(z.size(0) == y.size(0)) # refines i0 = s0 |
| if z.size(0) == 4: |
| return y * 2 |
| else: |
| return y + 2 |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.tensor(10), torch.randn(10)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, x_1, y_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(x_1); x_1 = None |
| zeros = torch.ops.aten.zeros.default([_local_scalar_dense], device = device(type='cpu'), pin_memory = False); _local_scalar_dense = zeros = None |
| add = torch.ops.aten.add.Tensor(y_1, 2); y_1 = None |
| return add""") # noqa: B950 |
| |
| @unittest.skipIf(not HAS_CUDA, 'CUDA-only test') |
| @unittest.expectedFailure |
| def test_unbacked_unify_guard_transitivity(self): |
| def f(x1, x2, y): |
| z1 = torch.zeros(x1.item()) |
| z2 = torch.zeros(x2.item()) |
| torch._check(z1.size(0) == z2.size(0)) # refines i0 = i1 |
| torch._check(z2.size(0) == y.size(0)) # refines i0 = s0 |
| if z1.size(0) == 4: |
| return y * 2 |
| else: |
| return y + 2 |
| |
| gm = make_fx(f, tracing_mode="symbolic")( |
| torch.tensor(10, device="cuda"), |
| torch.tensor(10, device="cuda"), |
| torch.randn(10, device="cuda") |
| ) |
| insert_deferred_runtime_asserts(gm, gm.shape_env, "test") |
| gm.recompile() |
| r = str(gm.code).strip() |
| # self.assertExpectedInline( |
| # r, """""" # noqa: B950 |
| # ) |
| |
| @unittest.skipIf(not HAS_CUDA, 'CUDA-only test') |
| def test_unbacked_unify_dependency_violation(self): |
| def f(x1, x2, x3, y): |
| z1 = x1.item() |
| torch._check(z1 // 9 == 1) |
| z2 = x2.item() |
| z3 = x3.item() |
| torch._check(z1 == z2 + z3) |
| return y * 2 |
| if z2 + z3 == z1: |
| return y * 2 |
| else: |
| return y + 3 |
| |
| # NB: inputs are done as CUDA to ensure they aren't queried to be |
| # backed |
| |
| gm = make_fx(f, tracing_mode="symbolic")( |
| torch.tensor(10, device="cuda"), torch.tensor(5, device="cuda"), |
| torch.tensor(5, device="cuda"), torch.randn(1, device="cuda") |
| ) |
| insert_deferred_runtime_asserts(gm, gm.shape_env, "test") |
| gm.recompile() |
| self.assertEqual(gm( |
| torch.tensor(12, device="cuda"), torch.tensor(6, device="cuda"), |
| torch.tensor(6, device="cuda"), torch.tensor([1.0], device="cuda")), |
| torch.tensor([2.0], device="cuda") |
| ) |
| with self.assertRaises(RuntimeError): |
| gm( |
| torch.tensor(20, device="cuda"), torch.tensor(10, device="cuda"), |
| torch.tensor(10, device="cuda"), torch.tensor([1.0], device="cuda") |
| ) |
| |
| |
| def test_split_unbacked_sizes(self): |
| def f(lengths, values): |
| # tolist not directly supported atm |
| sizes = [lengths[i].item() for i in range(lengths.size(0))] |
| for s in sizes: |
| # TODO(avik): no assertion generated with torch._check_is_size? |
| torch._constrain_as_size(s) |
| return torch.split(values, sizes) |
| |
| r = str(make_fx(f, tracing_mode="symbolic")( |
| torch.tensor([2, 3, 4]), |
| torch.randn(9) |
| ).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, lengths_1, values_1): |
| select = torch.ops.aten.select.int(lengths_1, 0, 0) |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(select); select = None |
| select_1 = torch.ops.aten.select.int(lengths_1, 0, 1) |
| _local_scalar_dense_1 = torch.ops.aten._local_scalar_dense.default(select_1); select_1 = None |
| select_2 = torch.ops.aten.select.int(lengths_1, 0, 2); lengths_1 = None |
| _local_scalar_dense_2 = torch.ops.aten._local_scalar_dense.default(select_2); select_2 = None |
| sym_constrain_range_for_size = torch.ops.aten.sym_constrain_range_for_size.default(_local_scalar_dense); sym_constrain_range_for_size = None |
| sym_constrain_range_for_size_1 = torch.ops.aten.sym_constrain_range_for_size.default(_local_scalar_dense_1); sym_constrain_range_for_size_1 = None |
| sym_constrain_range_for_size_2 = torch.ops.aten.sym_constrain_range_for_size.default(_local_scalar_dense_2); sym_constrain_range_for_size_2 = None |
| split_with_sizes = torch.ops.aten.split_with_sizes.default(values_1, [_local_scalar_dense, _local_scalar_dense_1, _local_scalar_dense_2]); values_1 = _local_scalar_dense = _local_scalar_dense_1 = _local_scalar_dense_2 = None |
| getitem = split_with_sizes[0] |
| getitem_1 = split_with_sizes[1] |
| getitem_2 = split_with_sizes[2]; split_with_sizes = None |
| return (getitem, getitem_1, getitem_2)""") # noqa: B950 |
| |
| def test_invalidate_nonzero(self): |
| ok = False |
| |
| def f(a): |
| nonlocal ok |
| b = a.clone() |
| x = b.nonzero() |
| x1 = b.nonzero() |
| x2 = b.nonzero() |
| assert x1.shape[0] == x2.shape[0] |
| ok = True |
| b.normal_() |
| y = b.nonzero() |
| try: |
| bool(x1.shape[0] == y.shape[0]) |
| self.fail("didn't raise exception") |
| except GuardOnDataDependentSymNode: |
| pass |
| |
| make_fx(f, tracing_mode="symbolic")(torch.randn(4)) |
| |
| @torch._functorch.config.patch(fake_tensor_propagate_real_tensors=True) |
| def test_invalidate_nonzero_propagate_real_tensors(self): |
| def f(a): |
| b = a.clone() |
| x = b.nonzero() |
| x1 = b.nonzero() |
| x2 = b.nonzero() |
| assert x1.shape[0] == x2.shape[0] |
| b.normal_() |
| y = b.nonzero() |
| # Because you're not actually going to generate exactly zero with |
| # normal_ lol |
| assert x1.shape[0] == y.shape[0] |
| |
| make_fx(f, tracing_mode="symbolic")(torch.randn(4)) |
| |
| def test_sqrt_size(self): |
| def f(a): |
| return a / a.size(-1) ** 0.5 |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(4)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| sym_size_int = torch.ops.aten.sym_size.int(a_1, 0) |
| sym_float = torch.sym_float(sym_size_int); sym_size_int = None |
| pow_1 = sym_float ** 0.5; sym_float = None |
| div = torch.ops.aten.div.Tensor(a_1, pow_1); a_1 = pow_1 = None |
| return div""") |
| |
| def test_make_fx_with_custom_tracer_preserving_nn_module_stack(self): |
| |
| class Bar(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| |
| def forward(self, x): |
| return x + 1 |
| |
| class Foo(torch.nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.bar = Bar() |
| |
| def forward(self, x): |
| return x + self.bar(x) |
| |
| gm = make_fx(Foo())(torch.randn(4, 4)) |
| for node in gm.graph.nodes: |
| self.assertTrue("nn_module_stack" not in node.meta) |
| |
| foo = Foo() |
| |
| def functional_call(*args, **kwargs): |
| with stateless._reparametrize_module(foo, {}): |
| return foo(*args, **kwargs) |
| |
| functional_call._orig_mod = foo |
| |
| gm_with_stack = make_fx(functional_call, record_module_stack=True)(torch.randn(4, 4)) |
| found = False |
| for node in gm_with_stack.graph.nodes: |
| if "nn_module_stack" in node.meta: |
| if len(node.meta["nn_module_stack"]) == 1: |
| self.assertTrue("custom_tracer_preserving_nn_module_stack.<locals>.Foo" in str(node.meta["nn_module_stack"])) |
| found = True |
| elif len(node.meta["nn_module_stack"]) == 2: |
| self.assertTrue("preserving_nn_module_stack.<locals>.Bar" in str(node.meta["nn_module_stack"])) |
| found = True |
| else: |
| # there can be at most 2 level |
| self.assertTrue(False) |
| |
| self.assertTrue(found) |
| |
| gm_without_stack = make_fx(functional_call)(torch.randn(4, 4)) |
| for node in gm_without_stack.graph.nodes: |
| self.assertTrue("nn_module_stack" not in node.meta) |
| |
| def test_symint_to_tensor(self): |
| def f(a): |
| return a / a.shape[0] |
| |
| r = str(make_fx(f, tracing_mode="symbolic")(torch.empty(4)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| sym_size_int = torch.ops.aten.sym_size.int(a_1, 0) |
| div = torch.ops.aten.div.Tensor(a_1, sym_size_int); a_1 = sym_size_int = None |
| return div""") |
| |
| r = str(make_fx(f, tracing_mode="symbolic", decomposition_table=decomposition_table)(torch.empty(4)).code).strip() |
| self.assertExpectedInline(r, """\ |
| def forward(self, a_1): |
| sym_size_int = torch.ops.aten.sym_size.int(a_1, 0) |
| sym_float = torch.sym_float(sym_size_int); sym_size_int = None |
| div = torch.ops.prims.div.default(a_1, sym_float); a_1 = sym_float = None |
| return div""") |
| |
| def test_cat(self): |
| def f(a, b): |
| val = torch.mul(a, b) |
| out = torch.cat([val, val]) |
| if out.shape[0] * out.shape[1] > 20: |
| out = out.cos() |
| return out |
| |
| test_inputs = [] |
| test_inputs.append([(1, 5), (6, 1)]) |
| test_inputs.append([(1, 4), (3, 1)]) |
| gm = self._test_dynamic(f, [(1, 6), (8, 1)], test_inputs) |
| self.assertTrue(eval_guards(gm, torch.randn(1, 10), torch.randn(6, 1))) |
| self.assertFalse(eval_guards(gm, torch.randn(1, 2), torch.randn(4, 1))) |
| self.assertExpectedInline(show_guards(gm), """2*L['a'].size()[1]*L['b'].size()[0] > 20""") |
| |
| def test_new_empty(self): |
| def f(a, b): |
| return a.new_empty(b.shape[0], b.shape[1] * 2) |
| |
| self._test_dynamic(f, [(2, 4), (4, 5)], [[(2, 3), (5, 7)], [(3, 7), (9, 3)]], assert_eq=False).shape_env |
| |
| def test_size_with_tensor(self): |
| # I think I messed up writing this test case originally, I think |
| # I'm supposed to hit an error case, but the code here works in both |
| # eager and tracing |
| def f(tensor): |
| max_size = torch.tensor([800, 1216], dtype=torch.int64) |
| batch_shape = [2] + list(tensor.shape[:-2]) + list(max_size) |
| return tensor.new_empty(batch_shape) |
| |
| a = torch.randn(3, 800, 1199) |
| f(a) |
| make_fx(f, tracing_mode="symbolic")(a) |
| |
| def test_fake_tensor_as_size(self): |
| def f(x): |
| r = torch.zeros([x]) |
| return r |
| |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.tensor(4)) |
| self.assertExpectedInline(fx_g.code.strip(), """\ |
| def forward(self, x_1): |
| _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(x_1); x_1 = None |
| zeros = torch.ops.aten.zeros.default([_local_scalar_dense], device = device(type='cpu'), pin_memory = False); _local_scalar_dense = None |
| return zeros""") # noqa: B950 |
| |
| def test_expand(self): |
| def f(a): |
| b = torch.mul(a, a) |
| c = b.expand(a.shape) |
| return c |
| |
| self._test_dynamic(f, [(3,)], [[(3,)], [(4,)], [(2,)]]) |
| self._test_dynamic(f, [(5, 1)], [[(4, 1)], [(3, 1)], [(6, 1)]]) |
| |
| def test_metadata(self): |
| def f(a, b): |
| d = a.new_empty(a.shape[0] + b.shape[0]) |
| return d |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.randn(5), torch.randn(4)) |
| meta_c = _get_node(fx_g, lambda x: x.target == aten.new_empty.default) |
| meta_d = _get_node(fx_g, lambda x: x.target == operator.add) |
| self.assertTrue(meta_c.meta['val'].shape[0].node.expr == meta_d.meta['val'].node.expr) |
| |
| def test_metadata_fresh(self): |
| def f(x): |
| assert x.shape[0] == 3 |
| return x.cos() |
| |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.randn(3)) |
| meta_cos = _get_node(fx_g, lambda x: x.target == aten.cos.default) |
| meta_inp = _get_node(fx_g, lambda x: x.op == 'placeholder') |
| self.assertTrue(meta_cos.meta['val'].shape[0] == 3) |
| # Checks if the input expr has been updated even though the constraint |
| # happened afterwards |
| self.assertTrue(meta_inp.meta['val'].shape[0] == 3) |
| |
| def test_elementwise_meta_with_sym_numbers(self): |
| def f(x, offset, as_sym_float=False): |
| x0 = x.size()[0] |
| if as_sym_float: |
| x0 = torch.sym_float(x0) |
| return torch.add(x0, offset) |
| |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.rand(2, 3), 2.0, False) |
| meta_add = _get_node(fx_g, lambda x: x.target == aten.add.Tensor) |
| self.assertEqual(meta_add.meta['val'].shape, ()) |
| self.assertEqual(meta_add.meta['val'].dtype, torch.float32) |
| |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.rand(2, 3), 2, False) |
| meta_add = _get_node(fx_g, lambda x: x.target == aten.add.Tensor) |
| self.assertEqual(meta_add.meta['val'].shape, ()) |
| self.assertEqual(meta_add.meta['val'].dtype, torch.int64) |
| |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.rand(2, 3), 2, True) |
| meta_add = _get_node(fx_g, lambda x: x.target == aten.add.Tensor) |
| self.assertEqual(meta_add.meta['val'].shape, ()) |
| self.assertEqual(meta_add.meta['val'].dtype, torch.float32) |
| |
| def test_return_symint(self): |
| def f(x): |
| return x.shape[0], x.cos(), x.shape[0] / 5 |
| self._test_dynamic(f, [(5,)], [[(4,)], [(12,)]]) |
| |
| def f(x): |
| return x.shape |
| self._test_dynamic(f, [(5, 3)], [[(4, 6)]]) |
| |
| def test_rmethod(self): |
| def f(x): |
| return x.size(0) + x |
| self._test_dynamic(f, [(5,)], [[(4,)], [(12,)]]) |
| |
| def test_mega_guard(self): |
| def f(a, b): |
| assert a.shape[0] == b.shape[0] * 2 |
| return a.cos() |
| fx_g = make_fx(f, tracing_mode="symbolic")(torch.randn(16), torch.randn(8)) |
| from torch._dynamo.source import LocalSource |
| self.assertExpectedInline( |
| str(fx_g.shape_env.produce_guards(fx_placeholder_vals(fx_g), [LocalSource("a"), LocalSource("b")], ignore_static=False)), # noqa: B950 |
| """["L['a'].size()[0] == 2*L['b'].size()[0]", "L['a'].stride()[0] == 1", "L['a'].storage_offset() == 0", "L['b'].stride()[0] == 1", "L['b'].storage_offset() == 0", "2 <= L['b'].size()[0]"]""" # noqa: B950 |
| ) |
| self.assertExpectedInline( |
| str(fx_g.shape_env.produce_guards(fx_placeholder_vals(fx_g), [LocalSource("a"), LocalSource("b")], ignore_static=True)), # noqa: B950 |
| """["L['a'].size()[0] == 2*L['b'].size()[0]", "2 <= L['b'].size()[0]"]""" # noqa: B950 |
| ) |
| |
| def test_guard_upperbound_range_refinement(self): |
| def f(a): |
| assert a.shape[0] > 5 and a.shape[0] > 12 |
| return a.cos() |
| tensor = make_fx(f, tracing_mode="symbolic")(torch.randn(15)) |
| self.assertExpectedInline(show_guards(tensor), """13 <= L['a'].size()[0]""") |
| |
| def test_guard_lowerbound_range_refinement(self): |
| def f(a): |
| assert a.shape[0] < 20 and a.shape[0] < 30 |
| return a.cos() |
| tensor = make_fx(f, tracing_mode="symbolic")(torch.randn(15)) |
| self.assertExpectedInline(show_guards(tensor), """L['a'].size()[0] <= 19""") |
| |
| def test_guard_upperbound_range_refinement_multivariate(self): |
| def f(a): |
| assert a.shape[0] > 5 and a.shape[0] > 12 |
| assert a.shape[1] > 5 and a.shape[1] > a.shape[0] |
| return a.cos() |
| tensor = make_fx(f, tracing_mode="symbolic")(torch.randn((15, 20))) |
| self.assertExpectedInline(show_guards(tensor), """\ |
| L['a'].size()[1] > L['a'].size()[0] |
| 13 <= L['a'].size()[0] |
| 14 <= L['a'].size()[1]""") |
| |
| def test_guard_lowerbound_range_refinement_multivariate(self): |
| def f(a): |
| assert a.shape[0] < 20 and a.shape[0] < 30 |
| assert a.shape[1] < 30 and a.shape[1] < a.shape[0] |
| return a.cos() |
| tensor = make_fx(f, tracing_mode="symbolic")(torch.randn((15, 5))) |
| self.assertExpectedInline( |
| show_guards(tensor), |
| """\ |
| L['a'].size()[1] < L['a'].size()[0] |
| L['a'].size()[0] <= 19 |
| L['a'].size()[1] <= 18""") |
| |
| def test_sym_storage_offset(self): |
| def f(x, y): |
| return x + y |
| |
| inp = (torch.randn(8)[3:], torch.randn(5)) |
| fx_g = make_fx(f, tracing_mode="symbolic")(*inp) |
| inp = (torch.randn(8)[3:], torch.randn(5)) |
| self.assertEqual(fx_g(*inp), f(*inp)) |
| |
| def _assert_no_guards(self, fx_g, free_symbols): |
| assert _get_free_symbols(fx_g.shape_env) == free_symbols, fx_g.shape_env.var_to_val |
| assert len(fx_g.shape_env.get_nontrivial_guards()) == 0, fx_g.shape_env.format_guards() |
| |
| def test_guards_equal(self): |
| def f(a, b): |
| return a * b |
| |
| # NB: Numbers are carefully chosen to avoid duck shaping from applying |
| |
| fx_g = _trace(f, (5, 6), (5, 6)) |
| self._assert_no_guards(fx_g, 2) |
| |
| fx_g = _trace(f, (5, 6, 7), (5, 6, 7)) |
| self._assert_no_guards(fx_g, 3) |
| |
| fx_g = _trace(f, (5, 1), (1, 6)) |
| self._assert_no_guards(fx_g, 2) |
| |
| def f(a, b, c, d): |
| a = a + b |
| cat = torch.cat([c, d]) |
| return a + cat |
| |
| fx_g = _trace(f, 7, 7, 4, 3) |
| self._assert_no_guards(fx_g, 2) |
| |
| def f(a, b, c, d, e): |
| vals = [a, b, c, d, e] |
| x = a |
| for idx in range(len(vals) - 1): |
| x = torch.cat([x, vals[idx]]) + vals[idx + 1] |
| return x |
| |
| fx_g = _trace(f, 2, 4, 8, 16, 32) |
| self._assert_no_guards(fx_g, 1) |
| |
| def f(a, b): |
| a = a.view(b.shape[0]) |
| return a + b.sum() |
| |
| fx_g = _trace(f, (4, 2), 8) |
| self._assert_no_guards(fx_g, 2) |
| |
| fx_g = _trace(f, (4, 2), (8, 5)) |
| self._assert_no_guards(fx_g, 3) |
| |
| fx_g = _trace(f, (2, 3, 4), 24) |
| self._assert_no_guards(fx_g, 3) |
| |
| def test_nonidentity_transitive_guards(self): |
| def f(a, b, c, d, e): |
| vals = [a, b, c, d, e] |
| cat_vals = [] |
| for idx in range(len(vals) - 1): |
| cat_vals.append(torch.cat([vals[idx], vals[idx]])) |
| final_vals = [] |
| for a, b in reversed(list(zip(cat_vals, vals[1:]))): |
| final_vals.append(a + b) |
| return final_vals |
| |
| fx_g = _trace(f, 2, 4, 8, 16, 32) |
| self.assertExpectedInline(show_guards(fx_g), """""") |
| |
| @torch.fx.experimental._config.patch(translation_validation=True) |
| def test_constant_specialization(self): |
| def f(t): |
| assert t.shape[0] == 10 |
| return t |
| |
| tensor = make_fx(f, tracing_mode="symbolic")(torch.randn(10)) |
| self.assertExpectedInline(show_guards(tensor), """""") |
| |
| |
| make_fx_failures = { |
| # unknown |
| xfail('allclose'), |
| xfail('equal'), |
| # empty |
| skip('new_empty'), |
| skip('empty_like'), |
| skip('empty'), |
| skip('empty_permuted'), |
| # flaky |
| skip('linalg.lstsq', 'grad_oriented'), |
| skip('nn.functional.max_unpool1d', '', device_type='cpu'), |
| skip('nn.functional.max_unpool2d', '', device_type='cpu'), |
| skip('nn.functional.max_unpool3d', '', device_type='cpu'), |
| skip('linalg.lstsq'), # flaky, probably just a precision issue |
| |
| # data-dependent control flow |
| skip('item'), |
| xfail('cov'), |
| xfail('nn.functional.gaussian_nll_loss'), |
| xfail('tensor_split'), |
| xfail('corrcoef'), |
| xfail('quantile'), |
| xfail('nanquantile'), |
| |
| # Seems like it's creating a sparse tensor that isn't captured by tensor.is_sparse |
| xfail('sparse.sampled_addmm'), |
| xfail('sparse.mm', 'reduce'), |
| |
| # proxy tensor doesn't support sparse correctly right now |
| skip('to_sparse'), |
| # segfaults |
| skip('block_diag'), |
| |
| # AssertionError: Tensor-likes are not close! |
| skip('empty_strided', '', device_type='cpu'), |
| } |
| |
| only_real_tensor_failures = { |
| xfail('narrow'), |
| } |
| |
| only_fake_tensor_failures = { |
| xfail('narrow'), |
| } |
| |
| fake_tensor_failures = { |
| # ASAN failures due to divide by 0 |
| skip('nn.functional.nll_loss'), |
| } |
| |
| symbolic_tensor_failures = { |
| xfail('combinations', ''), |
| xfail('geqrf', ''), # aten.geqrf.default - couldn't find symbolic meta function/decomposition |
| xfail('histogram', ''), # Could not run 'aten::histogram.bin_ct' with arguments from the 'Meta' backend. This c... |
| xfail('histogramdd', ''), # aten._histogramdd_bin_edges.default - couldn't find symbolic meta function/decomposition |
| xfail('nanquantile', ''), # Could not run 'aten::equal' with arguments from the 'Meta' backend. |
| xfail('nn.functional.binary_cross_entropy', ''), # aten.new_empty.default - couldn't find symbolic meta function/decom... |
| xfail('nn.functional.cross_entropy', ''), # aten.size.default - couldn't find symbolic meta function/decomposition |
| xfail('nn.functional.ctc_loss'), # aten._ctc_loss.Tensor - couldn't find symbolic meta function/decomposition |
| xfail('quantile', ''), # Could not run 'aten::equal' with arguments from the 'Meta' backend. |
| xfail('unique_consecutive', ''), # aten.unique_consecutive.default - couldn't find symbolic meta function/decomposition |
| |
| xfail('max_pool2d_with_indices_backward', ''), # Expected a value of type 'List[int]' for argument 'kernel_size' but... |
| |
| # many complex operators incorrect striding, metadata |
| xfail('fft.fft', ''), |
| xfail('fft.hfft2', ''), |
| xfail('fft.hfft', ''), |
| xfail('fft.hfftn', ''), |
| xfail('fft.ifft', ''), |
| xfail('fft.ihfft2', ''), |
| xfail('fft.ihfft', ''), |
| xfail('fft.ihfftn', ''), |
| xfail('fft.ihfft2', ''), |
| xfail('fft.irfft2', ''), |
| xfail('fft.irfft', ''), |
| xfail('fft.irfftn', ''), |
| xfail('fft.rfft2', ''), |
| xfail('fft.rfft', ''), |
| xfail('fft.rfftn', ''), |
| xfail('stft', '') |
| } |
| symbolic_tensor_segfaults = { |
| skip('nn.functional.batch_norm') # Segfault?? |
| } |
| |
| symbolic_tensor_failures.update(symbolic_tensor_segfaults) |
| |
| inplace_symbolic_tensor_failures = { |
| # bugs |
| xfail('float_power', ''), # base given to float_power_ has dtype Float but the operation's result requires dtype Double |
| } |
| |
| out_symbolic_tensor_failures = { |
| # Cast error details: Unable to cast (...) to Tensor |
| # |
| # This happens because the test is set up to call the out variant using the `out` kwarg: |
| # torch._some_op(arg1, arg2, out=(out1, out2, out3)) |
| # |
| # However, this only works on torch ops, not aten ops. For `_batch_norm_with_update`, |
| # this fails because the op has no python bindings, so it doesn't support the `out` kwarg |
| # way of calling its out variant. |
| xfail('_batch_norm_with_update', ''), |
| xfail('_native_batch_norm_legit', ''), |
| xfail('angle', ''), |
| xfail('argmax', ''), |
| xfail('argmin', ''), |
| xfail('fft.fft2', ''), |
| xfail('fft.fftn', ''), |
| xfail('fft.ifft2', ''), |
| xfail('fft.ifftn', ''), |
| xfail('gather', ''), |
| xfail('linalg.pinv', ''), |
| xfail('linalg.pinv', 'hermitian'), |
| xfail('lu', ''), |
| xfail('scatter_add', ''), |
| xfail('scatter', ''), |
| xfail('take_along_dim', ''), |
| xfail('triangular_solve', ''), |
| |
| # SymIntArrayRef expected to contain only concrete |
| xfail('ones', ''), |
| xfail('randn', ''), |
| xfail('zeros', ''), |
| |
| # RuntimeError: Cannot call numel() on tensor with symbolic sizes/strides |
| xfail('index_reduce', 'prod'), |
| xfail('index_reduce', 'mean'), |
| xfail('index_reduce', 'amax'), |
| xfail('index_reduce', 'amin'), |
| } |
| |
| out_symbolic_tensor_segfaults = { |
| skip('nanmean', ''), |
| } |
| |
| out_symbolic_tensor_failures.update(out_symbolic_tensor_segfaults) |
| |
| # Copies inputs to inplace operations to avoid inplace modifications |
| # to leaves requiring gradient |
| def _get_safe_inplace(inplace_variant): |
| @functools.wraps(inplace_variant) |
| def _fn(t, *args, **kwargs): |
| return inplace_variant(t.clone(), *args, **kwargs) |
| |
| return _fn |
| |
| def _test_make_fx_helper(self, device, dtype, op, tracing_mode, inplace=False, out=False): |
| fn = _get_safe_inplace(op.get_inplace()) if inplace else op.op |
| sample_inputs_itr = op.sample_inputs(device, dtype, requires_grad=False) |
| |
| # Limit ourselves to first 100 inputs so symbolic tracing tests don't take too long |
| count = 100 |
| if out: |
| count = 5 |
| for sample_input in itertools.islice(sample_inputs_itr, count): |
| if inplace and sample_input.broadcasts_input: |
| continue |
| args = [sample_input.input] + list(sample_input.args) |
| kwargs = sample_input.kwargs |
| if out: |
| expected = fn(*args, **kwargs) |
| kwargs['out'] = expected |
| |
| try: |
| optests.make_fx_check(fn, args, kwargs, tracing_mode, self.assertEqual, |
| randomize_data=True) |
| except DynamicOutputShapeException: |
| self.skipTest("Dynamic output shape operation in trace") |
| |
| |
| def skipIfNameMatches(pattern): |
| """ |
| Decorator to skip a test if its name matches the given pattern. |
| """ |
| def decorator(test_func): |
| def wrapper(*args, **kwargs): |
| if re.match(pattern, test_func.__name__): |
| raise unittest.SkipTest(f"Test '{test_func.__name__}' skipped because its name matches the pattern '{pattern}'") |
| return test_func(*args, **kwargs) |
| return wrapper |
| return decorator |
| |
| # Auto functionalize shouldn't work with make_fx directly |
| filtered_hop_db = [op for op in hop_db if op.name != "auto_functionalize"] |
| |
| @unittest.skipIf(not torch._dynamo.is_dynamo_supported(), "Cond requires dynamo") |
| class TestProxyTensorOpInfo(TestCase): |
| @ops(op_db + filtered_hop_db + custom_op_db, allowed_dtypes=(torch.float,)) |
| @skipOps('TestProxyTensorOpInfo', 'test_make_fx_exhaustive', make_fx_failures.union(only_real_tensor_failures)) |
| def test_make_fx_exhaustive(self, device, dtype, op): |
| _test_make_fx_helper(self, device, dtype, op, "real") |
| |
| @ops(op_db + filtered_hop_db + custom_op_db, allowed_dtypes=(torch.float,)) |
| @skipOps('TestProxyTensorOpInfo', 'test_make_fx_fake_exhaustive', |
| make_fx_failures.union(fake_tensor_failures, only_fake_tensor_failures)) |
| def test_make_fx_fake_exhaustive(self, device, dtype, op): |
| _test_make_fx_helper(self, device, dtype, op, "fake") |
| |
| @ops(op_db + filtered_hop_db + custom_op_db, allowed_dtypes=(torch.float,)) |
| @skipOps('TestProxyTensorOpInfo', 'test_make_fx_symbolic_exhaustive', |
| make_fx_failures | fake_tensor_failures | symbolic_tensor_failures) |
| def test_make_fx_symbolic_exhaustive(self, device, dtype, op): |
| _test_make_fx_helper(self, device, dtype, op, "symbolic") |
| |
| @ops(op_db + custom_op_db, allowed_dtypes=(torch.float,)) |
| @skipOps('TestProxyTensorOpInfo', 'test_make_fx_symbolic_exhaustive_inplace', |
| make_fx_failures | fake_tensor_failures | symbolic_tensor_failures | inplace_symbolic_tensor_failures) |
| def test_make_fx_symbolic_exhaustive_inplace(self, device, dtype, op): |
| if not op.get_inplace(): |
| self.skipTest("No inplace variable for this op") |
| _test_make_fx_helper(self, device, dtype, op, "symbolic", inplace=True) |
| |
| @ops(op_db + custom_op_db, allowed_dtypes=(torch.float,)) |
| @skipOps('TestProxyTensorOpInfo', 'test_make_fx_symbolic_exhaustive_out', |
| make_fx_failures | fake_tensor_failures | symbolic_tensor_failures | out_symbolic_tensor_failures) |
| def test_make_fx_symbolic_exhaustive_out(self, device, dtype, op): |
| if not op.supports_out: |
| self.skipTest("Op doesn't support out") |
| _test_make_fx_helper(self, device, dtype, op, "symbolic", out=True) |
| |
| |
| only_for = ("cpu") |
| instantiate_device_type_tests(TestProxyTensorOpInfo, globals(), only_for=only_for) |
| |
| |
| if __name__ == '__main__': |
| run_tests() |
