# Owner(s): ["module: decompositions"] import itertools import torch import os import numpy as np from enum import Enum from torch.overrides import resolve_name from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten from torch.utils import _pytree as pytree from torch._subclasses.meta_utils import MetaConverter, assert_metadata_eq, is_sparse_any import torch.utils._python_dispatch from torch._dispatch.python import enable_python_dispatcher from torch._ops import OpOverload, OpOverloadPacket from torch.testing import make_tensor from torch.testing._internal.common_utils import unMarkDynamoStrictTest from torch.testing._internal.common_utils import ( TestCase, skipIfCrossRef, skipIfTorchDynamo, suppress_warnings, TEST_WITH_ASAN, TEST_WITH_TORCHDYNAMO, run_tests, dtype_abbrs, parametrize ) from torch.testing._internal.common_device_type import ( ops, instantiate_device_type_tests, onlyCUDA, onlyCPU, OpDTypes, ) from torch.testing._internal.common_methods_invocations import ( binary_ufuncs, op_db, foreach_unary_op_db, foreach_binary_op_db, foreach_pointwise_op_db, foreach_reduce_op_db, foreach_other_op_db) from torch.testing._internal.opinfo.core import S, SampleInput from torchgen.yaml_utils import YamlLoader from torchgen.model import OperatorName import copy import sys import yaml import atexit import re from collections import defaultdict from collections.abc import Iterable import unittest import warnings import weakref from functools import partial, wraps bf16 = torch.bfloat16 f64 = torch.float64 f32 = torch.float32 f16 = torch.float16 c32 = torch.complex32 c64 = torch.complex64 c128 = torch.complex128 i8 = torch.int8 i16 = torch.int16 i32 = torch.int32 i64 = torch.int64 b8 = torch.bool u8 = torch.uint8 u16 = torch.uint16 u32 = torch.uint32 u64 = torch.uint64 foreach_op_db = ( foreach_unary_op_db + foreach_binary_op_db + foreach_pointwise_op_db + foreach_reduce_op_db + foreach_other_op_db ) class TestMetaConverter(TestCase): def assertSameVersionCounter(self, m1, m2): # Cannot easily test m1 and m2 have same storage due to # lack of Storage bindings. Use version counter. vc = m1._version self.assertEqual(m2._version, vc) # Doing it this way ensures that we get VC bump even with leaves with torch.no_grad(): m1._base.add_(3) self.assertNotEqual(m1._version, vc) self.assertEqual(m2._version, m1._version) def assertMetadataMatches(self, m1, m2): assert_metadata_eq(self.assertEqual, m1, m2) def test_view_of_non_leaf(self): x = torch.randn(4, requires_grad=True) y = x.neg() z1 = y[:] z2 = y[:] to_meta = MetaConverter() m1 = to_meta(z1) m2 = to_meta(z2) # check the test is actually testing what it claims self.assertTrue(m1._is_view()) self.assertFalse(m1._base.is_leaf) self.assertIsNot(m1, m2) self.assertMetadataMatches(m1, z1) self.assertMetadataMatches(m2, z2) self.assertSameVersionCounter(m1, m2) def test_view_of_leaf(self): x = torch.randn(4, requires_grad=True) z1 = x[:] z2 = x[:] to_meta = MetaConverter() m1 = to_meta(z1) m2 = to_meta(z2) # check the test is actually testing what it claims self.assertTrue(m1._is_view()) self.assertTrue(m1._base.is_leaf) self.assertIsNot(m1, m2) self.assertMetadataMatches(m1, z1) self.assertMetadataMatches(m2, z2) self.assertSameVersionCounter(m1, m2) def test_view_of_view_of_leaf(self): x = torch.randn(8) y = x.view(2, 4) y.requires_grad = True z = y.view(2, 2, 2) to_meta = MetaConverter() mx = to_meta(x) mz = to_meta(z) self.assertFalse(z.is_leaf) self.assertMetadataMatches(mx, x) self.assertMetadataMatches(mz, z) def test_leaf(self): x = torch.randn(4, requires_grad=True) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertTrue(m.is_leaf) self.assertTrue(m.requires_grad) self.assertMetadataMatches(m, x) def test_non_leaf(self): x = torch.randn(4, requires_grad=True) y = x.neg() to_meta = MetaConverter() m = to_meta(y) # check the test is actually testing what it claims self.assertFalse(m.is_leaf) self.assertTrue(m.requires_grad) self.assertMetadataMatches(m, y) def test_requires_grad_false(self): x = torch.randn(4, requires_grad=False) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertFalse(m.requires_grad) self.assertMetadataMatches(m, x) def test_channels_last(self): x = torch.empty(2, 3, 4, 5, memory_format=torch.channels_last) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertTrue(m.is_leaf) self.assertMetadataMatches(m, x) def test_channels_last_leaf(self): x = torch.empty(2, 3, 4, 5, memory_format=torch.channels_last, requires_grad=True) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertTrue(m.requires_grad) self.assertTrue(m.is_leaf) self.assertMetadataMatches(m, x) def test_channels_last_non_leaf(self): x = torch.empty(2, 3, 4, 5, memory_format=torch.channels_last, requires_grad=True) y = x + 2 # sanity self.assertEqual(x.stride(), y.stride()) self.assertFalse(y.is_leaf) to_meta = MetaConverter() m = to_meta(y) # check the test is actually testing what it claims self.assertTrue(m.requires_grad) self.assertFalse(m.is_leaf) self.assertMetadataMatches(m, y) # Check that we can autograd with m as input without erroring; # see https://github.com/pytorch/pytorch/issues/87956 loss = m.sum() torch.autograd.grad(loss, m) def test_empty_strided_non_dense_leaf(self): x = torch.empty_strided((2, 2), (4, 2), requires_grad=True) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertTrue(m.requires_grad) self.assertTrue(m.is_leaf) self.assertMetadataMatches(m, x) def test_view_mutate(self): x = torch.zeros(4) y = x.view(2, 2) to_meta = MetaConverter() m = to_meta(y) y.add_(torch.randn(2, 2, requires_grad=True)) m.add_(torch.randn(2, 2, device='meta', requires_grad=True)) def test_non_leaf_torture(self): x = torch.empty(20, requires_grad=True) with torch.no_grad(): x.set_(x.storage(), 10, (2,), (2,)) to_meta = MetaConverter() m = to_meta(x) # check the test is actually testing what it claims self.assertTrue(m.requires_grad) self.assertTrue(m.is_leaf) self.assertMetadataMatches(m, x) # NB: complex stuff is not actually exercised right now because # we have a blanket exclusion for complex conversion def test_view_as_real(self): x = torch.randn(4, dtype=torch.complex64) y = torch.view_as_real(x) m = MetaConverter()(y) self.assertMetadataMatches(m, y) def test_complex_noncontiguous_bug(self): x = torch.randn((2, 2, 4, 9), dtype=torch.complex32)[:, 0, :, :] m = MetaConverter()(x) self.assertMetadataMatches(m, x) def test_view_as_complex(self): x = torch.randn((4, 2), dtype=torch.float32) y = torch.view_as_complex(x) m = MetaConverter()(y) self.assertMetadataMatches(m, y) def test_view_dtype(self): x = torch.randn(4, dtype=torch.float32) y = x.view(dtype=torch.int32) m = MetaConverter()(y) self.assertMetadataMatches(m, y) def test_imag(self): x = torch.randn(4, dtype=torch.complex64) y = x.imag m = MetaConverter()(y) self.assertMetadataMatches(m, y) def test_inplace_set_storage(self): x = torch.tensor([0, 1], dtype=torch.int64) storage = x.untyped_storage() ssize = storage.size() meta = torch.empty((), dtype=torch.int64) meta.set_(storage, 0, (), ()) self.assertEqual(storage.size(), ssize) @skipIfTorchDynamo("https://github.com/pytorch/torchdynamo/issues/1991") def test_weakref(self): x = torch.randn(4, 4, 4) m = MetaConverter() y = m(x) z = m(x) self.assertIs(y, z) self.assertEqual(len(m.tensor_memo), 1) self.assertEqual(len(m.storage_memo), 1) self.assertEqual(len(m.describer.lookup_tensor), 1) self.assertEqual(len(m.describer.lookup_storage), 1) del x # Entries from Tensor -> int get deallocated when the real tensor # disappears... self.assertEqual(len(m.describer.lookup_tensor), 0) self.assertEqual(len(m.describer.lookup_storage), 0) del y del z # ... but the int -> FakeTensor entries don't die until the fake # tensors themselves die (because the user may have held onto the # int key and are expecting to get a consistent fake tensor in # this case) self.assertEqual(len(m.tensor_memo), 0) self.assertEqual(len(m.storage_memo), 0) li = [] r = [] for i in range(4): li.append(torch.rand([i])) r.append(m(li[-1])) self.assertEqual(len(m.tensor_memo), 4) self.assertEqual(len(m.storage_memo), 4) self.assertEqual(len(m.describer.lookup_tensor), 4) self.assertEqual(len(m.describer.lookup_storage), 4) del li self.assertEqual(len(m.describer.lookup_tensor), 0) self.assertEqual(len(m.describer.lookup_storage), 0) del r self.assertEqual(len(m.tensor_memo), 0) self.assertEqual(len(m.storage_memo), 0) @skipIfTorchDynamo("https://github.com/pytorch/torchdynamo/issues/1991") def test_tensor_outlives_converter(self): m = MetaConverter() ref = weakref.ref(m) x = torch.randn([4, 4]) y = m(x) del m self.assertIs(ref(), None) aten = torch.ops.aten CHECK_STRIDES = { torch.Tensor.__getitem__, } CHECK_ALL_STRIDES = { aten.unsqueeze.default } CHECK_STRIDES_SKIPS = { aten._conj_physical.default, aten._fft_c2c.default, aten._fft_c2r.default, aten._fft_r2c.default, aten._linalg_svd.default, aten.binary_cross_entropy.default, aten.complex.default, aten.polar.default, aten.copysign.Tensor, aten.div.Tensor_mode, aten.floor_divide.default, aten.heaviside.default, aten.lerp.Scalar, aten.lerp.Tensor, aten.logaddexp.default, aten.logical_and.default, aten.logical_or.default, aten.logical_xor.default, aten.pow.Scalar, aten.prelu.default, aten.special_xlog1py.default, aten.xlogy.Tensor, aten.nll_loss2d_forward.default, # channel_last and channel_last_3d related failures aten.convolution.default, # following ops fails if include_storage_offset = True, but these are a bit edge casey # we should still fix them, leaving them here for tracking. # aten._reshape_alias.default, # repro with test_dispatch_symbolic_meta_outplace_all_strides_matmul_cuda_float32 # aten.view.default, # repro with test_dispatch_symbolic_meta_outplace_all_strides_unflatten_cuda_float32 } CHECK_CONJ_SKIPS = { # The conj bit is not copied, see: # https://github.com/pytorch/pytorch/pull/101836 aten.linalg_lu_solve.out, } class CheckStrides(Enum): NONE = 0 SIGNIFICANT = 1 ALL = 2 def should_check_strides(func): if func in CHECK_ALL_STRIDES: return CheckStrides.ALL if func in CHECK_STRIDES: return CheckStrides.SIGNIFICANT if func in CHECK_STRIDES_SKIPS: return CheckStrides.NONE if not isinstance(func, torch._ops.OpOverload): return CheckStrides.NONE # Prims are expected to model strides correctly if func.namespace == "prims": return CheckStrides.SIGNIFICANT # Check if it's a view, by testing if any of the returns have # a non-empty alias set if any(r.alias_info.before_set for r in func._schema.returns if r.alias_info): return CheckStrides.SIGNIFICANT # TODO: check for TensorIterator return CheckStrides.SIGNIFICANT def assert_ref_meta_equal(test_case, func, meta_rs, rs, msg_callable): flat_meta_rs = pytree.tree_leaves(meta_rs) flat_rs = pytree.tree_leaves(rs) test_case.assertEqual(len(flat_meta_rs), len(flat_rs)) for i, meta_r, r in zip(range(len(flat_rs)), flat_meta_rs, flat_rs): def test_assert(cond, msg): if not cond: raise RuntimeError(f"output {i}: {msg_callable(msg)}") if not isinstance(r, torch.Tensor): continue test_assert(isinstance(meta_r, torch.Tensor), f"but real {i}th result is Tensor") test_assert(meta_r.dtype == r.dtype, f"for element {i}, was {meta_r.dtype} but real dtype was {r.dtype}") test_assert(meta_r.shape == r.shape, f"for element {i}, was {meta_r.shape} but real shape was {r.shape}") # See https://github.com/pytorch/pytorch/issues/78050 if should_check_strides(func) == CheckStrides.ALL: same_strides, _ = torch._prims_common.check_all_strides(meta_r, r) test_assert(same_strides, f"for element {i}, was {meta_r.stride()} but real stride was {r.stride()}") elif should_check_strides(func) == CheckStrides.SIGNIFICANT: same_strides, _ = torch._prims_common.check_significant_strides(meta_r, r) test_assert(same_strides, f"for element {i}, was {meta_r.stride()} but real stride was {r.stride()}") test_assert( meta_r.storage_offset() == r.storage_offset(), f"for element {i}, was {meta_r.storage_offset()} but real storage_offset was {r.storage_offset()}") test_assert(meta_r.requires_grad == r.requires_grad, f"for element {i}, was {meta_r.requires_grad} but real requires_grad was {r.requires_grad}") if func not in CHECK_CONJ_SKIPS: test_assert(meta_r.is_conj() == r.is_conj(), f"for element {i}, was {meta_r.is_conj()} but real is_conj was {r.is_conj()}") test_assert(meta_r.is_neg() == r.is_neg(), f"for element {i}, was {meta_r.is_neg()} but real is_neg was {r.is_neg()}") # This environment variable controls whether or not we print expected failure # lists at the end of a test suite run. The intended usage looks like this: # # 1. Run `PYTORCH_COLLECT_EXPECT=1 python test/test_meta.py` on a CUDA build # of PyTorch that has LAPACK/MAGMA installed. You can filter `-k test_meta` # or `-k test_dispatch_meta` to only focus on one or another list # 2. Given the printed skip/xfail list, add them to the corresponding lists; # torch.* entries go in meta_function and aten.* entries go in meta_dispatch. # If there are preexisting entries, you need to merge in the entries. # # This is somewhat manual but typically you shouldn't need to do this, unless # you've made a major change (e.g., added a new dtype to PyTorch) and need to # refresh the lists. If you want to do it from scratch, just clear out the # preexisting lists before running. # # WARNING: Python dict literals will silently ignore duplicate keys COLLECT_EXPECT = os.getenv('PYTORCH_COLLECT_EXPECT', '0') == '1' seen_succeeded = {} seen_failed = {} failed_reasons = defaultdict(set) def print_seen(): expected_failures = [] skips = [] def fmt_dtypes(dtypes): r = ', '.join(sorted(dtype_abbrs[d] for d in dtypes)) return '{' + r + '}' for op, failed_dtypes in seen_failed.items(): ops = resolve_name(op) succeeded_dtypes = seen_succeeded.get(op, set()) expected_failures_dtypes = failed_dtypes - succeeded_dtypes skips_dtypes = failed_dtypes & succeeded_dtypes reasons = "" if failed_reasons[op]: reasons = " # " + ", ".join(sorted(failed_reasons[op])) if expected_failures_dtypes: expected_failures.append(f" {ops}: {fmt_dtypes(expected_failures_dtypes)},{reasons}") if skips_dtypes: skips.append(f" {ops}: {fmt_dtypes(skips_dtypes)},") expected_failures.sort() skips.sort() nl = '\n' print(f"""\ expected_failures = {{ {nl.join(expected_failures)} }} skips = {{ {nl.join(skips)} }} """) if COLLECT_EXPECT: atexit.register(print_seen) # Success forces pass; failure forces fail; skip unconditionally skips testing TestExpect = Enum("TestExpect", ("SUCCESS", "XFAILURE", "SKIP")) # unlike print produce strides def verbose_print(e): class Lit: def __init__(self, s): self.s = s def __repr__(self): return self.s def go(t): if is_sparse_any(t): return t elif isinstance(t, torch.Tensor): return Lit(f"{t} stride={t.stride()}") else: return t return repr(tree_map(go, e)) def run_meta_crossref( test_case, test_expect, func, args, kwargs, *, dtype, device_type, run_symbolic_meta: bool ): to_meta = MetaConverter() do_meta = test_expect is not TestExpect.SKIP if do_meta: try: meta_args = tree_map(to_meta, args) meta_kwargs = tree_map(to_meta, kwargs) except Exception as e: raise RuntimeError( f"failed to convert args to meta; " f"originally (*{args}, **{kwargs})") from e try: rs = func(*args, **kwargs) except Exception as e: raise AssertionError("Original OpInfo is broken") from e # TODO: also handle cases where func raise an exception # For now, only attempt if we managed to convert all tensor types # (if any of them failed, we're in a mixed device situation and # this isn't well supported) if do_meta and to_meta.successful(): # Special cases if func is torch.tensor_split: # Use original indices_or_sections, this argument is data dependent meta_args = (meta_args[0], args[1]) + meta_args[2:] elif func is torch.Tensor.__getitem__: # Ensure boolean tensors use original assert len(args) == 2 flat_args = pytree.tree_leaves(args[1]) flat_meta_args, spec = tree_flatten(meta_args[1]) flat_new_args = [] for a, ma in zip(flat_args, flat_meta_args): flat_new_args.append(a if isinstance(a, torch.Tensor) and a.dtype in [torch.int8, torch.bool] else ma) meta_args = (meta_args[0], tree_unflatten(flat_new_args, spec)) elif func in (torch.ops.aten.repeat_interleave.Tensor, torch.ops.aten.repeat_interleave.Tensor_out): if kwargs.get("output_size", None) is None: meta_args = args if func is torch.ops.aten.repeat_interleave.Tensor_out: meta_kwargs["out"] = kwargs["out"] elif func in (torch.ops.aten.index.Tensor, torch.ops.aten.index.Tensor_out): # Don't convert boolean tensors to meta as they will have nonzero # called on them indices = [] for meta_index, real_index in zip(meta_args[1], args[1]): if meta_index is not None and meta_index.dtype in [torch.int8, torch.bool]: indices.append(real_index) else: indices.append(meta_index) meta_args = (meta_args[0], indices) elif func is torch.nn.functional.ctc_loss and all([isinstance(args[2], list), isinstance(args[3], list)]): # torch.ops.aten._ctc_loss.IntList has a meta kernel but # torch.ops.aten._ctc_loss.Tensor does not test_expect = TestExpect.SUCCESS if kwargs.get("device", None) is not None: meta_kwargs["device"] = "meta" try: # Suppress warnings, this doesn't matter for test_meta.py # but it does matter if you want to use this decorator # for cross-ref testing, as some tests may be looking at # errors with warnings.catch_warnings(): warnings.simplefilter("ignore") if run_symbolic_meta: # Run the decomps and meta kernels registered # to the python dispatcher instead of the regular dispatcher. # This should be the same set of kernels # that fake tensor runs in dynamic shapes mode. with enable_python_dispatcher(): meta_rs = func(*meta_args, **meta_kwargs) else: meta_rs = func(*meta_args, **meta_kwargs) except Exception as e: if test_expect is TestExpect.XFAILURE: return rs seen_failed.setdefault(func, set()).add(dtype) if isinstance(e, NotImplementedError): m = RE_NOT_IMPLEMENTED_MSG.search(e.args[0]) if m: failed_reasons[func].add(m.group(1)) if COLLECT_EXPECT: return rs raise RuntimeError(f"""\ failed to run: {resolve_name(func)}( *{verbose_print(meta_args)}, **{verbose_print(meta_kwargs)} )""") from e else: try: delim = ',\n ' assert_ref_meta_equal(test_case, func, meta_rs, rs, lambda msg: f"""\ meta disagrees with real impl: {resolve_name(func)}( {delim.join(map(verbose_print, meta_args))}, {delim.join(k + ": " + verbose_print(v) for k, v in meta_kwargs.items())} ) = ( {verbose_print(meta_rs)} ) {msg} """) except Exception: if test_expect is TestExpect.XFAILURE: return rs seen_failed.setdefault(func, set()).add(dtype) if COLLECT_EXPECT: return rs raise else: seen_succeeded.setdefault(func, set()).add(dtype) if test_expect is TestExpect.XFAILURE and not COLLECT_EXPECT: raise RuntimeError(f"unexpected success {resolve_name(func)} {meta_args} {meta_kwargs}") return rs RE_NOT_IMPLEMENTED_MSG = re.compile(r"Could not run '([^']+)' with arguments ") meta_function_expected_failures = { torch.Tensor.to_sparse : {f64, i32, c128, i64, i16, f16, u8, c64, bf16, b8, i8, f32}, torch.allclose : {f64, f16, c128, c64, bf16, f32}, torch.argwhere : {f64, i32, c128, i64, i16, f16, u8, c64, bf16, b8, i8, f32}, torch.combinations : {f64, i32, c128, i64, i16, f16, u8, c64, bf16, b8, i8, f32}, torch.corrcoef : {f64, i32, c128, i64, i16, u8, c64, bf16, f16, i8, f32}, torch.cov : {f64, i32, c128, i64, i16, u8, c64, bf16, i8, f32, f16}, torch.functional.istft : {f64, c64, c128, f32}, torch.geqrf : {f64, c64, c128, f32}, torch.masked_select : {f64, i32, c128, i64, i16, f16, u8, c64, bf16, b8, i8, f32}, torch.nonzero : {f64, i32, c128, i64, i16, c32, f16, u8, c64, bf16, b8, i8, f32}, torch.Tensor.nonzero : {f64, i32, c128, i64, i16, c32, f16, u8, c64, bf16, b8, i8, f32}, torch.Tensor.item : {f64, i32, c128, i64, i16, f16, u8, c32, c64, bf16, b8, i8, f32}, torch.bincount : {i32, i64, u8, i16, i8}, torch.functional.unique : {f64, i32, i64, u8, i16, f16, bf16, b8, i8, f32, u16, u32, u64}, torch.functional.unique_consecutive : {f64, i32, i64, u8, i16, f16, bf16, b8, i8, f32, u16, u32, u64}, torch.histogram : {f64, f32}, torch.histogramdd : {f64, f32}, torch.nn.functional.ctc_loss : {f64, f32}, torch.nn.functional.gaussian_nll_loss : {f16, f64, bf16, f32}, torch.linalg.lstsq : {f64, f32, c128, c64}, } meta_function_expected_failures_conditional = { torch.repeat_interleave : (lambda dtype, *args, **kwargs: not isinstance(kwargs.get("repeats", None), int)), } """ # This is some sample code for how we could dump these dicts into YAML # file for easier reading/writing import yaml print(yaml.dump( {resolve_name(k): [dtype_abbrs[d] for d in v] for k, v in meta_function_expected_failures.items()}, default_flow_style=None)) import sys sys.exit() """ meta_function_skips = { torch.Tensor.__rmatmul__ : {bf16, c128, f64, f32, f16, c64}, torch.Tensor.matmul : {f64, f32, c128, c64}, torch.functional.atleast_2d : {bf16, i8, c32, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.functional.atleast_3d : {bf16, i8, c32, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.functional.cartesian_prod : {bf16, i8, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.functional.einsum : {bf16, c128, f64, f32, f16, c64}, torch.inner : {f16, bf16, i8, i64, u8, c128, f64, i16, f32, i32, c64}, torch.linalg.matrix_norm : {c128, f32, c64, f64}, torch.linalg.matrix_rank : {c128, c64}, torch.linalg.svd : {c128, c64}, torch.matmul : {bf16, c128, f64, f32, f16, c64}, torch.nanquantile : {f64, f32}, torch.narrow : {bf16, i8, i64, u8, c128, b8, f64, i16, i32, f32, f16, c32, c64}, torch.nn.functional.batch_norm : {f64, f32}, torch.nn.functional.binary_cross_entropy : {bf16, f64, f32, f16}, torch.nn.functional.dropout3d : {bf16, f64, f32, f16}, torch.nn.functional.local_response_norm : {bf16, f64, f32, f16}, torch.svd : {c128, c64}, torch.take_along_dim : {bf16, i8, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.vstack : {bf16, i8, c32, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.diff : {b8}, torch.equal : {bf16, i8, c32, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.nanmean : {bf16, f64, f32, f16, c32, c64, c128}, torch.nn.functional.cross_entropy : {bf16, f64, f32}, torch.nn.functional.nll_loss : {bf16, f64, f32}, torch.linalg.cond : {c128, c64, f32, f64}, torch.linalg.vecdot : {bf16, f64, f32, f16}, torch.empty : {bf16, i8, c32, i64, u8, c128, b8, f64, i16, i32, f32, f16, c64}, torch.Tensor.addbmm_: {bf16, c128, c64, f32, f64, i16, i32, i64, i8, u8}, torch.nn.functional.one_hot : {i64}, } meta_function_device_expected_failures = defaultdict(dict) meta_function_device_expected_failures_only_outplace = defaultdict(dict) meta_function_device_skips = defaultdict(dict) meta_function_device_expected_failures['cpu'] = { # TODO: The decomps for these batch norm ops return different dtypes depending # on the device. We should make this work better with meta tensors. torch.native_batch_norm: {bf16, f16}, torch._native_batch_norm_legit: {bf16, f16}, torch.ops.aten._batch_norm_with_update: {bf16, f16}, torch.native_layer_norm: {bf16, f16}, } meta_function_device_expected_failures['cuda'] = { torch.corrcoef: {bf16, f16}, # aten::_local_scalar_dense torch.cov: {f16}, # aten::_local_scalar_dense torch.functional.unique: {f16}, # aten::_unique2, aten::unique_dim torch.functional.unique_consecutive: {f16}, # aten::unique_consecutive torch.geqrf: {f32, f64}, # aten::geqrf } meta_function_device_skips['cpu'] = { # TODO: The decomps for these batch norm ops return different dtypes depending # on the device. We should make this work better with meta tensors. torch.native_batch_norm: {f32, f64}, torch._native_batch_norm_legit: {f32, f64}, torch.ops.aten._batch_norm_with_update: {f32, f64}, } meta_function_device_skips['cuda'] = { torch.inner: {f16}, torch.linalg.matrix_rank: {f32, f64}, torch.linalg.svd: {f32, f64}, torch.nn.functional.cross_entropy: {f16}, torch.nn.functional.interpolate: {f16}, torch.nn.functional.nll_loss: {f16}, torch.svd: {f32, f64}, } # This is a __torch_function__ mode that, when enabled, interposes every # Torch API call and runs the operator as normal, and then reruns it # with meta inputs, and then checks that everything about the output agrees. # Most of the logic deals with faithfully replicating the original tensor # as a meta tensor, which is nontrivial because there are a lot of subsystems # that may potentially be exercised. # # That being said, this class is a little overkill for what it is doing in # this test file (since I could have just inlined __torch_function__ on the # OpInfo call, and OpInfos generally have very regular inputs), but it will be # useful for more comprehensive testing e.g., as seen in # https://github.com/pytorch/pytorch/pull/75994 The big benefit is it is # A LOT more efficient that torch dispatch mode (at the cost of less coverage) class MetaCrossRefFunctionMode(torch.overrides.TorchFunctionMode): test_case: TestCase device_type: str dtype: torch.dtype def __init__(self, test_case, *, device, dtype, inplace): self.test_case = test_case self.device_type = torch.device(device).type self.dtype = dtype self.inplace = inplace def __torch_function__(self, func, types, args=(), kwargs=None): kwargs = kwargs or {} if ( torch.jit.is_tracing() or isinstance(func, torch.ScriptMethod) or # meta converter doesn't work correctly when no_dispatch() is on, so # skip running the crossref test in this case torch._C._dispatch_tls_local_exclude_set().has(torch._C.DispatchKey.Python) ): return func(*args, **kwargs) if self.dtype in meta_function_skips.get(func, set()): test_expect = TestExpect.SKIP elif self.dtype in meta_function_device_skips[self.device_type].get(func, set()): test_expect = TestExpect.SKIP elif self.dtype in meta_function_expected_failures.get(func, set()): test_expect = TestExpect.XFAILURE elif self.dtype in meta_function_device_expected_failures[self.device_type].get(func, set()): test_expect = TestExpect.XFAILURE elif meta_function_expected_failures_conditional.get(func, lambda *_, **__: False)(self.dtype, *args, **kwargs): test_expect = TestExpect.XFAILURE elif not self.inplace and \ self.dtype in meta_function_device_expected_failures_only_outplace[self.device_type].get(func, set()): test_expect = TestExpect.XFAILURE else: test_expect = TestExpect.SUCCESS return run_meta_crossref( self.test_case, test_expect, func, args, kwargs, dtype=self.dtype, device_type=self.device_type, run_symbolic_meta=False ) # these always fail meta_dispatch_expected_failures = { aten.allclose.default: {f16, bf16, f32, f64, c64, c128}, # NotImplementedError: 'aten::_local_scalar_dense' aten.geqrf.default : {c64, c128, f64, f32}, aten.linalg_lstsq.default : {c64, c128, f64, f32}, aten.masked_select.default : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten.masked_select.out : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten.nonzero.default : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, c32, b8, i16, u8}, aten.nonzero.out : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, c32, b8, i16, u8}, aten._to_sparse.default : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten._to_sparse.sparse_dim : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten._ctc_loss.Tensor : {f32, f64}, # Shape of second output depends on data. aten._histogramdd_bin_edges.default : {f32, f64}, aten._histogramdd_from_bin_cts.default : {f32, f64}, aten._histogramdd_from_bin_tensors.default : {f32, f64}, aten._local_scalar_dense.default : {c32, c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten._unique2.default : {i8, f64, i64, f16, bf16, f32, i32, b8, i16, u8, u16, u32, u64}, aten.bincount.default : {i64, i8, i32, i16, u8}, aten.equal.default : {c64, f16, i8, f64, c128, i64, bf16, f32, i32, b8, i16, u8}, aten.histogram.bin_ct : {f32, f64}, aten.histogram.bins_tensor : {f32, f64}, aten.unique_consecutive.default : {i8, f64, i64, f16, bf16, f32, i32, b8, i16, u8, u16, u32, u64}, aten.unique_dim.default : {i8, f64, i64, f16, bf16, f32, i32, b8, i16, u8, u16, u32, u64}, aten.upsample_nearest3d.vec : {bf16, f32, f64, u8}, } # these sometimes pass and sometimes fail meta_dispatch_skips = { aten.index.Tensor: {i64, bf16, f16, u8, b8, f32, i8, f64, i16, i32, c32, c64, c128}, # at::nonzero doesn't have a Meta function aten._to_copy.default: {i64, bf16, f16, u8, b8, f32, i8, f64, i16, i32, c32, c64, c128}, aten.empty.memory_format: {b8, bf16, c128, c64, c32, f16, f32, f64, i16, i32, i64, i8, u8}, aten.addbmm_.default: {bf16, c128, c64, f32, f64, i16, i32, i64, i8, u8}, } # For CompositeImplicitAutograd functions that fail before hitting the Mode meta_dispatch_early_skips = set({ torch.Tensor.float_power_, # Errors out in one of the tests, while ProxyTensor passes... torch.Tensor.cumprod_, torch.Tensor.cumsum_, }) meta_inplace_skips = set({ # Errors out in one of the tests, while ProxyTensor passes... torch.Tensor.cumprod_, torch.Tensor.cumsum_, }) meta_dispatch_device_expected_failures = defaultdict(dict) meta_dispatch_device_skips = defaultdict(dict) meta_dispatch_device_expected_failures['cpu'] = { # TODO: The decomps for these batch norm ops return different dtypes depending # on the device. We should make this work better with meta tensors. aten.native_batch_norm.default: {bf16, f16}, aten._native_batch_norm_legit.default: {bf16, f16}, aten._native_batch_norm_legit.no_stats: {bf16, f16}, aten._batch_norm_with_update.default: {bf16, f16}, aten.native_layer_norm.default: {bf16, f16}, } meta_dispatch_device_expected_failures['cuda'] = { aten._unique2.default: {f16}, # aten::_unique2 aten._use_cudnn_ctc_loss.default: {f32, f64}, # aten::_use_cudnn_ctc_loss aten._use_cudnn_ctc_loss.Tensor: {f32, f64}, # aten::_use_cudnn_ctc_loss.Tensor aten.cudnn_grid_sampler.default: {f16, f32, f64}, # aten::cudnn_grid_sampler aten.geqrf.default: {f32, f64}, # aten::geqrf aten.linalg_eigvalsh.out: {f32, f64}, # aten::linalg_eigvalsh.out aten.log_sigmoid_forward.default: {bf16, f16, f64, f32}, aten.log_sigmoid_forward.output : {bf16, f16, f64, f32}, # aten::log_sigmoid_forward.output aten.unique_consecutive.default: {f16}, # aten::unique_consecutive aten.unique_dim.default: {f16}, # aten::unique_dim aten.upsample_nearest3d.vec: {f16}, # aten::upsample_nearest3d.vec } meta_dispatch_device_skips['cpu'] = { aten._embedding_bag_forward_only.default: {bf16, f16, f32, f64}, # TODO: The decomps for these batch norm ops return different dtypes depending # on the device. We should make this work better with meta tensors. aten.native_batch_norm.default: {f32, f64}, aten._native_batch_norm_legit.default: {f32, f64}, aten._native_batch_norm_legit.no_stats: {f32, f64}, aten._batch_norm_with_update.default: {f32, f64}, # If the computation dtype is different from the input # dtype this will fail. CPU execution may also have a # a different output from other devices. aten.native_batch_norm.out: {bf16, f16, f32, f64} } meta_dispatch_device_skips['cuda'] = { aten._conj.default: {c32, f16}, # file issue aten._linalg_svd.default: {c64, c128}, # aten::linalg_eigvalsh.out aten.cudnn_batch_norm.default: {f32, f64}, aten.log_softmax.int : {c32, c64}, aten.softmax.int : {c32, c64}, aten.softmax.int : {c32, c64}, # ROCm stuff; technically this should be expected failure but it's # not worth it; these should get unified anyway aten.miopen_batch_norm.default: {f32}, } def get_strided_args(args): def get_strided_variants(t, include_storage_offset=False): variants = [] # contiguous variants.append(t) # transposed if t.ndim > 1: perm = list(reversed(range(t.ndim))) transposed = torch.empty( t.shape[::-1], device=t.device, dtype=t.dtype, requires_grad=t.requires_grad ).permute(perm).copy_(t) variants.append(transposed) # nondense if t.ndim > 0: nondense = torch.repeat_interleave(t, 2, dim=-1)[..., ::2] variants.append(nondense) # channel_last if t.ndim == 4: variants.append(t.contiguous(memory_format=torch.channels_last)) # channel_last_3d if t.ndim == 5: variants.append(t.contiguous(memory_format=torch.channels_last_3d)) # storage_offset if include_storage_offset: buffer = torch.empty(t.numel() + 1, device=t.device, dtype=t.dtype, requires_grad=t.requires_grad) buffer = buffer.as_strided(t.shape, t.stride(), storage_offset=1) buffer.copy_(t) variants.append(buffer) return variants strided_args = [] for arg in args: if isinstance(arg, torch.Tensor) and not arg.is_sparse_csr and arg.is_contiguous(): strided_arg_variants = get_strided_variants(arg) else: strided_arg_variants = [arg] strided_args.append(strided_arg_variants) yield from itertools.product(*strided_args) class MetaCrossRefDispatchMode(torch.utils._python_dispatch.TorchDispatchMode): test_case: TestCase device: torch.device dtype: torch.dtype aten_olp_no_out_overload: set = set() def __init__(self, test_case, *, device, dtype, symbolic_meta: bool, inplace: bool, supports_out: bool): self.test_case = test_case # save TLS self.precision = test_case.precision self.rel_tol = test_case.rel_tol self.device_type = torch.device(device).type self.dtype = dtype self.symbolic_meta = symbolic_meta self.inplace = inplace self.supports_out = supports_out @staticmethod def try_resolve_aten_out_overload(ol, args, kwargs, num_outputs): ol_args = ol._schema.arguments olp: OpOverloadPacket = ol._overloadpacket if olp in MetaCrossRefDispatchMode.aten_olp_no_out_overload: return (None, None, None) candidate_ols = [] for candidate_ol_name in olp.overloads(): candidate_ol = getattr(olp, candidate_ol_name) if any(arg.is_out for arg in candidate_ol._schema.arguments): candidate_ols.append(candidate_ol) if not candidate_ols: MetaCrossRefDispatchMode.aten_olp_no_out_overload.add(olp) return (None, None, None) # Now match based on args, kwargs and number of required outputs candidate_ol: OpOverload = None for candidate_ol in candidate_ols: candidate_ol_args = candidate_ol._schema.arguments if (len(args) >= len(candidate_ol_args)): continue # Positional arguments must have the same type if not all( ol_args[pos_arg_ind].type == candidate_ol_args[pos_arg_ind].type for pos_arg_ind in range(len(args)) ): continue # Number of outputs must match candidate_out_names = [out_arg.name for out_arg in candidate_ol_args[-num_outputs:] if out_arg.is_out] if len(candidate_out_names) != num_outputs: continue # Now try and match kwargs. Just need to ensure that the # remaining kwargs allow an out overload to be called. For example # we can throw away parameters like `dtype` that may be passed to the # functional version of the op since the `dtype` will already be present # in the `out` argument new_kwargs = {} kwargs_match = True for arg in candidate_ol_args[len(args):-num_outputs]: if arg.name not in kwargs: if arg.has_default_value(): new_kwargs[arg.name] = arg.default_value elif isinstance(arg.type, torch.OptionalType): if isinstance(arg.type.getElementType(), torch.BoolType): new_kwargs[arg.name] = False else: new_kwargs[arg.name] = None else: kwargs_match = False break else: new_kwargs[arg.name] = kwargs[arg.name] if kwargs_match: return candidate_ol, candidate_out_names, new_kwargs return None, None, None def _get_expected_test_result(self, func: OpOverload): if self.dtype in meta_dispatch_skips.get(func, set()): test_expect = TestExpect.SKIP elif self.dtype in meta_dispatch_device_skips[self.device_type].get(func, set()): test_expect = TestExpect.SKIP elif self.dtype in meta_dispatch_expected_failures.get(func, set()): test_expect = TestExpect.XFAILURE elif self.dtype in meta_dispatch_device_expected_failures[self.device_type].get(func, set()): test_expect = TestExpect.XFAILURE else: test_expect = TestExpect.SUCCESS return test_expect def __torch_dispatch__(self, func, types, args=(), kwargs=None): kwargs = kwargs or {} self.test_case.precision = self.precision self.test_case.rel_tol = self.rel_tol test_expect = self._get_expected_test_result(func) expected = run_meta_crossref( self.test_case, test_expect, func, args, kwargs, dtype=self.dtype, device_type=self.device_type, run_symbolic_meta=self.symbolic_meta, ) # This is to test torch ops that do not have an out parameter but have # aten op overloads that have out parameters. Additionally, Python decompositions # may register OpOverloadPacket's so decompositions need to be tested # to ensure all OpOverloads still function for the Meta key (e.g. if a python decomposition # is registered for an aten op aten.foo with overloads [default, out], the python # function needs to support receiving `out` arguments) if ( not self.inplace and not self.supports_out and test_expect == TestExpect.SUCCESS and (torch.is_tensor(expected) or isinstance(expected, Iterable)) ): # check to see if there is a potential out overload num_outputs = 1 if torch.is_tensor(expected) else len(expected) func_out_overload, out_param_names, kwargs = self.try_resolve_aten_out_overload(func, args, kwargs, num_outputs) if func_out_overload: if num_outputs == 1: kwargs[out_param_names[0]] = expected else: for ind, out_param_name in enumerate(out_param_names): kwargs[out_param_name] = expected[ind] test_expect = self._get_expected_test_result(func_out_overload) run_meta_crossref( self.test_case, test_expect, func_out_overload, args, kwargs, dtype=self.dtype, device_type=self.device_type, run_symbolic_meta=self.symbolic_meta, ) return expected # NB: we're running these tests only on CUDA because there are some # inconsistencies between CUDA and CPU, and running on CUDA makes it easier # to ignore the CPU case when inconsistencies arise. Ideally we deal # with the inconsistencies but this takes time. @unMarkDynamoStrictTest class TestMeta(TestCase): # Copies inputs to inplace operations to avoid inplace modifications # to leaves requiring gradient def _get_safe_inplace(self, inplace_variant): @wraps(inplace_variant) def _fn(t, *args, **kwargs): if isinstance(t, list): return inplace_variant([x.clone() for x in t], *args, **kwargs) else: return inplace_variant(t.clone(), *args, **kwargs) return _fn @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_meta_outplace(self, device, dtype, op): if "_scaled_mm" in op.name: raise unittest.SkipTest("_scaled_mm dose not support meta device") skip_op_names = ( "fft.ihfft", "fft.ihfft2", "linalg.lu_solve", ) if TEST_WITH_TORCHDYNAMO and op.name in skip_op_names: raise unittest.SkipTest("flaky") # run the OpInfo sample inputs, cross-referencing them with the # meta implementation and check the results are the same. All # the heavy lifting happens in MetaCrossRefFunctionMode func = op.get_op() samples = op.sample_inputs(device, dtype, requires_grad=False) for sample_input in samples: args = [sample_input.input] + list(sample_input.args) kwargs = sample_input.kwargs with MetaCrossRefFunctionMode(self, dtype=dtype, device=device, inplace=False): expected = func(*args, **kwargs) if isinstance(expected, torch.Tensor) and op.supports_out: func(*args, **kwargs, out=expected) # Special test for functions taking "device" kwarg # The crossref tests that replacing the device with "meta" works # This part makes sure that *_like functions work well with a "meta" # Tensor and their original device argument. if "device" in kwargs and "_like" in op.name: with torch.random.fork_rng(): torch.manual_seed(123) ref = func(*args, **kwargs) # *_like functions take a Tensor as first argument assert isinstance(args[0], torch.Tensor) with torch.random.fork_rng(): torch.manual_seed(123) args[0] = args[0].to(device="meta") meta = func(*args, **kwargs) # empty_like is not deterministic if op.name != "empty_like": self.assertEqual(ref, meta) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_meta_inplace(self, device, dtype, op): func = op.get_inplace() if not func: self.skipTest("No inplace variable for this op") if op.promotes_int_to_float and not dtype.is_floating_point: self.skipTest("Op promotes to float, which is impossible for inplace with non-float input") if func in meta_inplace_skips: self.skipTest("Skipped") func = self._get_safe_inplace(func) samples = op.sample_inputs(device, dtype, requires_grad=False) for sample_input in samples: if sample_input.broadcasts_input: continue args = [sample_input.input] + list(sample_input.args) kwargs = sample_input.kwargs with MetaCrossRefFunctionMode(self, dtype=dtype, device=device, inplace=True): expected = func(*args, **kwargs) def _run_dispatch_meta_test(self, device, dtype, op, symbolic_meta, inplace, all_stride_variants=False): if "_scaled_mm" in op.name: raise unittest.SkipTest("_scaled_mm dose not support meta device") if inplace: func = op.get_inplace() if not func: self.skipTest("No inplace variable for this op") if op.promotes_int_to_float and not dtype.is_floating_point: self.skipTest("Op promotes to float, which is impossible for inplace with non-float input") else: func = op.get_op() if func in meta_dispatch_early_skips: self.skipTest("Function is in dispatch early skips") if inplace: func = self._get_safe_inplace(func) samples = op.sample_inputs(device, dtype, requires_grad=False) for sample_input in samples: if inplace and sample_input.broadcasts_input: continue sample_args = [sample_input.input] + list(sample_input.args) kwargs = sample_input.kwargs if all_stride_variants and sum(isinstance(arg, torch.Tensor) for arg in sample_args) <= 5: # test inputs <= 5 tensors to avoid combinatorial explosion strided_args = get_strided_args(sample_args) else: strided_args = [sample_args] for args in strided_args: with MetaCrossRefDispatchMode.push( self, dtype=dtype, device=device, symbolic_meta=symbolic_meta, inplace=inplace, supports_out=op.supports_out): expected = func(*args, **kwargs) if not inplace and isinstance(expected, torch.Tensor) and op.supports_out: func(*args, **kwargs, out=expected) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_dispatch_meta_outplace(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=False, inplace=False) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_dispatch_meta_inplace(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=False, inplace=True) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_dispatch_symbolic_meta_outplace(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=True, inplace=False) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings @ops(itertools.chain(op_db, foreach_op_db)) def test_dispatch_symbolic_meta_inplace(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=True, inplace=True) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings # only test one dtype, as output stride behavior is the same for all dtypes @ops(itertools.chain(op_db, foreach_op_db), dtypes=OpDTypes.any_common_cpu_cuda_one) # Only test on CUDA, as CUDA kernel's stride is the reference @onlyCUDA def test_dispatch_symbolic_meta_outplace_all_strides(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=True, inplace=False, all_stride_variants=True) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings # only test one dtype, as output stride behavior is the same for all dtypes @ops(itertools.chain(op_db, foreach_op_db), dtypes=OpDTypes.any_common_cpu_cuda_one) # Only test on CUDA, as CUDA kernel's stride is the reference @onlyCUDA def test_dispatch_symbolic_meta_inplace_all_strides(self, device, dtype, op): self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=True, inplace=True, all_stride_variants=True) @unittest.skipIf(TEST_WITH_ASAN, "Skipped under ASAN") @skipIfCrossRef @suppress_warnings # only test one dtype, as output stride behavior is the same for all dtypes @ops(binary_ufuncs, allowed_dtypes=(torch.float32,)) # Only test on CUDA, as CUDA kernel's stride is the reference @onlyCUDA def test_binary_ufuncs_mixed_dtype(self, device, dtype, op): make_arg = partial( make_tensor, device=device, ) def sample_input(op, device, dtype, requires_grad, **kwargs): yield SampleInput( make_arg((S,), dtype=dtype), make_arg((S,), dtype=torch.float16) ) op = copy.copy(op) op.sample_inputs_func = sample_input self._run_dispatch_meta_test(device, dtype, op, symbolic_meta=True, inplace=False) def test_empty_quantized(self): r = torch.empty(2 ** 52, device='meta', dtype=torch.qint8) self.assertEqual(r.device.type, 'meta') def test_nan_to_num(self): t = torch.tensor([float('nan'), float('inf'), -float('inf'), 3.14], device='meta') r = t.nan_to_num() self.assertEqual(r.device.type, 'meta') def test_inplace_masked_fill_error(self): t = torch.randn(3, 3, device='meta') with self.assertRaisesRegex(RuntimeError, "doesn't match the broadcast"): t.masked_fill_((t > 0).unsqueeze(0), 0.1) def test_inplace_bin_ops_error(self): t = torch.randn(3, 3, device='meta') for op in (torch.Tensor.add_, torch.Tensor.sub_, torch.Tensor.mul_, torch.Tensor.div_, torch.Tensor.logical_and_, torch.Tensor.logical_or_, torch.Tensor.logical_xor_): with self.assertRaisesRegex(RuntimeError, "doesn't match the broadcast"): op(t, t.clone().unsqueeze(0)) @onlyCPU def test_meta_autograd_no_error(self): with torch.library._scoped_library("meta_test", "DEF") as lib: with torch.library._scoped_library("meta_test", "IMPL", "CPU") as impl_cpu: with torch.library._scoped_library("meta_test", "IMPL", "Meta") as impl_meta: def foo_impl(x): return x + 1 lib.define("foo(Tensor a) -> Tensor") impl_meta.impl("foo", foo_impl) impl_cpu.impl("foo", foo_impl) a = torch.ones(2, device='meta') # The point of the test is that this should not error: # We have a fallthrough kernel registered to the AutogradMeta # key for custom ops, so it's fine that `foo()` doesn't have # an autograd kernel. b = torch.ops.meta_test.foo.default(a) def test_huber_loss_backward(self): inps = [torch.rand(2**52, device='meta') for _ in range(3)] r = torch.ops.aten.huber_loss_backward(*inps, 0, 1.0) self.assertEqual(r.device.type, 'meta') self.assertEqual(r.shape, inps[0].shape) def _norm_backwards_test_helper(self, op, args, output_mask, expected_shapes): dtype = torch.float32 device = "meta" # test functional call grads = op(*args, output_mask) def assertEqualShapes(res, exp): self.assertIsNone(res) if exp is None else self.assertEqual(exp, res.shape) assertEqualShapes(grads[0], expected_shapes[0]) assertEqualShapes(grads[1], expected_shapes[1]) assertEqualShapes(grads[2], expected_shapes[2]) out_kwargs = { f"out{i}": torch.empty(0, device=device, dtype=dtype) for i in range(len(output_mask)) } # test call with out parameters grads = op(*args, output_mask, **out_kwargs) def assertEqualShapes(res, exp): self.assertEqual(exp, res.shape) if exp is not None else True assertEqualShapes(out_kwargs["out0"], expected_shapes[0]) assertEqualShapes(out_kwargs["out1"], expected_shapes[1]) assertEqualShapes(out_kwargs["out2"], expected_shapes[2]) @onlyCPU @parametrize("output_mask", list(itertools.product([True, False], [True, False], [True, False]))) def test_layer_norm_backward(self, output_mask): from torch.testing._internal.common_methods_invocations import sample_inputs_layer_norm device = "meta" dtype = torch.float32 samples = sample_inputs_layer_norm(None, device, dtype, requires_grad=False) for sample in samples: with self.subTest(sample=sample): # handle optional weight and bias if len(sample.args) != 3: sample.args = (*sample.args, *([None] * (3 - len(sample.args)))) grad_out = torch.ones_like(sample.input) normalized_shape, weight, bias = sample.args ndims_after_reduction = sample.input.ndim - len(normalized_shape) mean_shape = grad_out.shape[:ndims_after_reduction] mean = torch.zeros(mean_shape, device=device, dtype=dtype) rstd = torch.zeros(mean_shape, device=device, dtype=dtype) expected_shapes = ( sample.input.shape if output_mask[0] else None, weight.shape if output_mask[1] and weight is not None else None, bias.shape if output_mask[2] and bias is not None else None) args = [grad_out, sample.input, normalized_shape, mean, rstd, weight, bias] self._norm_backwards_test_helper(torch.ops.aten.native_layer_norm_backward, args, output_mask, expected_shapes) @onlyCPU @parametrize("output_mask", list(itertools.product([True, False], [True, False], [True, False]))) def test_group_norm_backward(self, output_mask): from torch.testing._internal.common_methods_invocations import sample_inputs_group_norm # input, (args) num_groups, (kwargs) weight, bias eps device = "meta" dtype = torch.float32 samples = sample_inputs_group_norm(None, device, dtype, requires_grad=False) for sample in samples: with self.subTest(sample=sample): grad_out = torch.ones_like(sample.input) N, C = sample.input.shape[:2] HxW = torch.prod(torch.as_tensor(sample.input.shape[2:]), dtype=torch.int32).item() group = sample.args[0] mean = torch.zeros((N, group), device=device, dtype=dtype) rstd = torch.zeros((N, group), device=device, dtype=dtype) weight = torch.zeros((C), device=device, dtype=dtype) args = [grad_out, sample.input, mean, rstd, weight, N, C, HxW, group] expected_shapes = ( sample.input.shape if output_mask[0] else None, weight.shape if output_mask[1] else None, weight.shape if output_mask[2] else None) # test functional call self._norm_backwards_test_helper(torch.ops.aten.native_group_norm_backward, args, output_mask, expected_shapes) @onlyCPU @parametrize("output_mask", list(itertools.product([True], [True, False], [True, False]))) def test_batch_norm_backward(self, output_mask): from torch.testing._internal.common_methods_invocations import sample_inputs_batch_norm # input, (args) num_groups, (kwargs) weight, bias eps device = "meta" dtype = torch.float32 samples = sample_inputs_batch_norm(None, device, dtype, requires_grad=False) for sample in samples: with self.subTest(sample=sample): if sample.input.dim() < 2: continue grad_out = torch.ones_like(sample.input) running_mean, running_var, weight, bias = sample.args train = sample.kwargs.get("training", True) save_mean = torch.zeros((sample.input.shape[1], ), device=device, dtype=dtype) if train else None save_invstd = torch.zeros((sample.input.shape[1], ), device=device, dtype=dtype) if train else None args = [grad_out, sample.input, weight, running_mean, running_var, save_mean, save_invstd, train, sample.kwargs.get("eps", 1e-5)] expected_shapes = ( sample.input.shape, torch.Size([sample.input.shape[1]]) if output_mask[1] else None, torch.Size([sample.input.shape[1]]) if output_mask[2] else None) self._norm_backwards_test_helper(torch.ops.aten.native_batch_norm_backward, args, output_mask, expected_shapes) def test_fill__alias_relationship(self): inps = torch.rand(2**52, device='meta') r = torch.ops.aten.fill_(inps, 1.0) # aten.fill_ returns an aliase self.assertEqual(id(inps), id(r)) # aten.fill returns a new tensor r2 = torch.ops.aten.fill(inps, 1.0) self.assertNotEqual(id(inps), id(r2)) def test_meta__fused_moving_avg_obs_fq_helper(self, device): from torch.ao.quantization import FusedMovingAvgObsFakeQuantize to_meta = MetaConverter() x = torch.randn(5, 5, device=device) running_min_op = torch.tensor(float("inf"), device=device) running_max_op = torch.tensor(float("-inf"), device=device) avg_const = 0.01 scale = torch.tensor([1.0], device=device) zero_point = torch.tensor([0], dtype=torch.int, device=device) mod = FusedMovingAvgObsFakeQuantize() torch.ao.quantization.enable_fake_quant(mod) torch.ao.quantization.enable_observer(mod) mod.to(device) meta_x = to_meta(x) args = [ x, mod.observer_enabled, mod.fake_quant_enabled, running_min_op, running_max_op, scale, zero_point, avg_const, 0, 255, 0, ] meta_args = args.copy() meta_args[0] = meta_x kwargss = [ {}, {"per_row_fake_quant": False, "symmetric_quant": False}, {"per_row_fake_quant": False, "symmetric_quant": True}, ] for kwargs in kwargss: ref_out = aten._fused_moving_avg_obs_fq_helper.default(*args, **kwargs) meta_out = aten._fused_moving_avg_obs_fq_helper.default(*meta_args, **kwargs) self.assertEqual(ref_out[0].size(), meta_out[0].size()) self.assertEqual(ref_out[0].stride(), meta_out[0].stride()) self.assertEqual(ref_out[1].size(), meta_out[1].size()) self.assertEqual(ref_out[1].stride(), meta_out[1].stride()) def test_cdist_forward(self, device): to_meta = MetaConverter() x1 = torch.rand([3, 2], device=device) x2 = torch.rand([2, 2], device=device) p = 2.0 for compute_mode in (None, 1, 2): ref = aten._cdist_forward.default(x1, x2, p, compute_mode) res = aten._cdist_forward.default(to_meta(x1), to_meta(x2), p, compute_mode) self.assertEqual(res.device.type, 'meta') self.assertEqual(ref.shape, res.shape) def test_quantized_embedding_bag(self): tab_shape = [8, 128] emb_size, ind_len, off_len = tab_shape[0], 32, 33 f_table = torch.from_numpy((np.random.random_sample(tab_shape) + 1).astype(np.float32)) q_table = torch.ops.quantized.embedding_bag_byte_prepack(f_table) indices = torch.from_numpy(np.random.randint(low=0, high=emb_size, size=ind_len)).int() max_length = len(indices) // (off_len - 1) if max_length > 20: max_length = 20 np_lengths = np.random.randint(0, max_length + 1, size=off_len - 1).astype(np.int32) offsets = torch.cat([torch.zeros([1]), torch.cumsum(torch.from_numpy(np_lengths), 0)]).int() eb = torch.ops.quantized.embedding_bag_byte_rowwise_offsets( q_table.to(device="meta"), indices.to(device="meta"), offsets.to(device="meta"), mode=0, # sum per_sample_weights=None, include_last_offset=True, ) self.assertEqual(eb.shape, [32, 128]) self.assertEqual(eb.dtype, torch.float32) self.assertEqual(eb.untyped_storage().data_ptr(), 0) # Tests mean and max. # Can't easily test sum, because there is a fast path for sum which # causes offset2bag to not get allocated... but the backward function # needs it, and the offset2bag computation lives inside the # derivatives.yaml formula directly, so there is no way to access it. # To test sum, need to manually compute offset2bag @parametrize("mode", [1, 2]) def test_embedding_bag_dense_backward(self, mode): weight = torch.randn(4, 3, requires_grad=True) indices = torch.tensor([1, 0, 2, 1, 3]) offsets = torch.tensor([0, 2, 3, 5]) scale_grad_by_freq = False sparse = False per_sample_weights = None include_last_offset = False padding_idx = -1 output, offset2bag, bag_size, maximum_indices = torch.ops.aten._embedding_bag.default( weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx ) grad = torch.randn_like(output) # Call the function with example inputs grad_weight = torch.ops.aten._embedding_bag_dense_backward.default( grad, indices, offset2bag, bag_size, maximum_indices, weight.size(0), scale_grad_by_freq, mode, per_sample_weights, padding_idx ) meta_grad_weight = torch.ops.aten._embedding_bag_dense_backward.default( grad.to('meta'), indices.to('meta'), offset2bag.to('meta'), bag_size.to('meta'), maximum_indices.to('meta'), weight.size(0), scale_grad_by_freq, mode, per_sample_weights, padding_idx ) self.assertEqual(grad_weight.to('meta'), meta_grad_weight) def test_embedding_bag_dense_backward_per_sample_weights(self): weight = torch.randn(4, 3, requires_grad=True) indices = torch.tensor([1, 0, 2, 1, 3]) offsets = torch.tensor([0, 2, 3, 5]) scale_grad_by_freq = False sparse = False mode = 0 per_sample_weights = torch.randn(5, requires_grad=True) include_last_offset = False padding_idx = -1 output, offset2bag, bag_size, maximum_indices = torch.ops.aten._embedding_bag.default( weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx ) grad = torch.randn_like(output) # Call the function with example inputs grad_weight = torch.ops.aten._embedding_bag_per_sample_weights_backward.default( grad, weight, indices, offsets, offset2bag, mode, padding_idx ) meta_grad_weight = torch.ops.aten._embedding_bag_per_sample_weights_backward.default( grad.to('meta'), weight.to('meta'), indices.to('meta'), offsets.to('meta'), offset2bag.to('meta'), mode, padding_idx ) self.assertEqual(grad_weight.to('meta'), meta_grad_weight) # opinfo test is using aten.fill_, it's not testing aten.fill @onlyCUDA def test_fill_stride(self): to_meta = MetaConverter() sample_args = [torch.rand(2, 2, 2, 2), 1.0] for args in get_strided_args(sample_args): meta_args = to_meta(args) ref_out = torch.ops.aten.fill(*args) meta_out = torch.ops.aten.fill(*meta_args) self.assertEqual(ref_out.size(), meta_out.size()) self.assertEqual(ref_out.stride(), meta_out.stride()) def test_map_location_deserialize(self): import io t = torch.rand(10) b = io.BytesIO() torch.save(t, b) b.seek(0) r = torch.load(b, map_location=torch.device("meta")) self.assertEqual(r.device.type, 'meta') self.assertEqual(r.shape, t.shape) self.assertEqual(r.dtype, t.dtype) self.assertEqual(r.storage().data_ptr(), 0) def test_embedding_bag_byte_prepack(self): batch_size = 10 num_embeddings = 80 embedding_dim = [128, 256, 512] res_shape = [[batch_size, num_embeddings, ed + 8] for ed in embedding_dim] for ed, rs in zip(embedding_dim, res_shape): weight = torch.randn(batch_size, num_embeddings, ed, dtype=torch.float32) res = torch.ops.quantized.embedding_bag_byte_prepack(weight.to(device="meta")) self.assertEqual(res.shape, rs) self.assertEqual(res.dtype, torch.float32) self.assertEqual(res.untyped_storage().data_ptr(), 0) def test_embedding_bag_byte_unpack(self): batch_size = 10 num_embeddings = 80 embedding_dim = [128, 256, 512] res_shape = [[batch_size, num_embeddings, ed] for ed in embedding_dim] for ed, rs in zip(embedding_dim, res_shape): packed_weight = torch.randn(batch_size, num_embeddings, ed + 8, dtype=torch.float32) res = torch.ops.quantized.embedding_bag_byte_unpack(packed_weight.to(device="meta")) self.assertEqual(res.shape, rs) self.assertEqual(res.dtype, torch.float32) self.assertEqual(res.untyped_storage().data_ptr(), 0) def test_index_select_out(self): def f(): input = torch.randn([8, 16], device='meta') index = torch.tensor([2, 1, 6, 7, 3, 1, 7, 5, 6, 7], device='meta') out = torch.empty([10, 16], device='meta') return torch.index_select(input=input, dim=0, index=index, out=out) with enable_python_dispatcher(): out = f() self.assertEqual(out.shape, [10, 16]) def test_local_scalar_dense_call(self): with self.assertRaisesRegex(RuntimeError, "cannot be called on meta tensors"): meta_tensor = torch.randn(1, device='meta') meta_tensor.item() instantiate_device_type_tests(TestMeta, globals()) def print_op_str_if_not_supported(op_str): op = OperatorName.parse(op_str) packet = getattr(torch.ops.aten, str(op.name)) overload = getattr(packet, op.overload_name if op.overload_name else "default") if any(overload in d for d in [meta_dispatch_skips, meta_dispatch_device_skips['cuda']]): print(f"{overload} # SKIP") if any(overload in d for d in [meta_dispatch_expected_failures, meta_dispatch_device_expected_failures['cuda']]): print(overload) if __name__ == "__main__": COMPARE_XLA = os.getenv('PYTORCH_COMPARE_XLA', None) if COMPARE_XLA is not None: with open(COMPARE_XLA) as f: d = yaml.load(f, Loader=YamlLoader) ops = d.get("full_codegen", []) + d.get("supported", []) + d.get("autograd", []) for op_str in ops: print_op_str_if_not_supported(op_str) sys.exit(0) COMPARE_TEXT = os.getenv('PYTORCH_COMPARE_TEXT', None) if COMPARE_TEXT is not None: with open(COMPARE_TEXT) as f: for op_str in f: print_op_str_if_not_supported(op_str.strip()) sys.exit(0) run_tests()