# Owner(s): ["module: functorch"] # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import itertools import unittest from torch.testing._internal.common_utils import TestCase, run_tests, is_iterable_of_tensors, IS_ARM64, parametrize import torch from torch import Tensor import functools from torch.testing._internal.common_device_type import instantiate_device_type_tests from torch.testing._internal.common_device_type import ops from torch.testing._internal.common_device_type import \ toleranceOverride, tol from functorch_additional_op_db import additional_op_db from torch.testing._internal.common_methods_invocations import op_db from common_utils import ( get_fallback_and_vmap_exhaustive, generate_vmap_inputs, decorate, xfail, skip, skipOps, tol1, tol2, opsToleranceOverride, check_vmap_fallback, is_batch_norm_training, is_valid_inplace_sample_input, loop2, ) from torch.testing._internal.opinfo.core import SampleInput from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map from functorch import grad, vjp, vmap, jacrev, jacfwd import torch.autograd.forward_ad as fwAD from functorch._src.eager_transforms import _as_tuple, jvp aten = torch.ops.aten # Version of autograd.grad with some differences: # - pytree inputs is allowed (but leaves of the pytree have to all # be tensors) # - if an input is not used as part of derivatives, we will return a # zero-filled tensor for the result def _autograd_grad( outputs, inputs, grad_outputs=None, retain_graph=False, create_graph=True ): inputs, inputs_spec = tree_flatten(inputs) diff_inputs = tuple(inp for inp in inputs if inp.requires_grad) if grad_outputs is None: diff_outputs = tuple(out for out in outputs if out.requires_grad) else: diff_grad_outputs = [ (out, go) for out, go in zip(outputs, grad_outputs) if out.requires_grad ] if len(diff_grad_outputs) == 0: diff_outputs, grad_outputs = (), () else: diff_outputs, grad_outputs = zip(*diff_grad_outputs) grad_inputs = torch.autograd.grad( diff_outputs, diff_inputs, grad_outputs, retain_graph=retain_graph, create_graph=create_graph, allow_unused=True, ) result = [] grad_inputs_iter = iter(grad_inputs) for inp in inputs: if inp.requires_grad: grad_input = next(grad_inputs_iter) if grad_input is None: result.append(torch.zeros_like(inp)) else: result.append(grad_input) else: result.append(torch.zeros_like(inp)) return tree_unflatten(result, inputs_spec) def diff_arg(arg, requires_grad=True): def is_differentiable_arg(arg): if requires_grad: return arg.requires_grad else: return arg.is_floating_point() or arg.is_complex() if is_iterable_of_tensors(arg): if all([is_differentiable_arg(a) for a in arg]): return True if all([not is_differentiable_arg(a) for a in arg]): return False raise RuntimeError("NYI: The test runner can't handle this") return isinstance(arg, Tensor) and is_differentiable_arg(arg) # Given f, returns an f' such that: # - f' takes only positional arguments # - All arguments to f' are floating-point Tensors # - All outputs of f' are floating-point Tensors def normalize_op_input_output2(f, args, kwargs, output_process_fn_grad=None, requires_grad=True): flat_args, args_spec = tree_flatten(args) diff_argnums = tuple(i for i, arg in enumerate(flat_args) if diff_arg(arg, requires_grad=requires_grad)) assert len(diff_argnums) > 0 primals = tuple(flat_args[i] for i in diff_argnums) @functools.wraps(f) def wrapped(*primals): _args = list(flat_args) for num, arg in zip(diff_argnums, primals): _args[num] = arg _args = tree_unflatten(_args, args_spec) result = f(*_args, **kwargs) if output_process_fn_grad is not None: result = output_process_fn_grad(result) if isinstance(result, tuple): result = tuple(r for r in result if torch.is_floating_point(r)) assert len(result) > 0 return result return wrapped, primals # TODO: consolidate with normalize_op_input_output2 def normalize_op_input_output3(f, args, kwargs, sample_args, output_process_fn_grad=None): flat_args, args_spec = tree_flatten(args) flat_sample_args, _ = tree_flatten(sample_args) diff_argnums = tuple(i for i, (arg, sample) in enumerate(zip(flat_args, flat_sample_args)) if diff_arg(sample, requires_grad=True)) assert len(diff_argnums) > 0 primals = tuple(flat_args[i] for i in diff_argnums) @functools.wraps(f) def wrapped(*primals): _args = list(flat_args) for num, arg in zip(diff_argnums, primals): _args[num] = arg _args = tree_unflatten(_args, args_spec) result = f(*_args, **kwargs) if output_process_fn_grad is not None: result = output_process_fn_grad(result) if isinstance(result, tuple): result = tuple(r for r in result if torch.is_floating_point(r)) assert len(result) > 0 return result return wrapped, primals def normalize_op_input_output(f, sample, requires_grad=True): args = tuple([sample.input] + list(sample.args)) return normalize_op_input_output2( f, args, sample.kwargs, sample.output_process_fn_grad, requires_grad=requires_grad ) def ref_vjp(f, *primals): result = f(*primals) def wrapped(cotangents): return _autograd_grad(_as_tuple(result), primals, _as_tuple(cotangents)) return result, wrapped def simulate_jvp(f, primals, tangents): primals_out, tangents_out = torch.autograd.functional.jvp(f, primals, tangents) return primals_out, tangents_out def ref_jvp(f, primals, tangents): with fwAD.dual_level(): duals = tuple(fwAD.make_dual(p, t) for p, t in zip(primals, tangents)) result_duals = f(*duals) result_duals, spec = tree_flatten(result_duals) primals_out, tangents_out = zip(*(fwAD.unpack_dual(d) for d in result_duals)) return tree_unflatten(primals_out, spec), tree_unflatten(tangents_out, spec) def get_sample_cotangents(f, sample): fn, primals = normalize_op_input_output(f, sample) output = fn(*primals) return tree_map(torch.randn_like, output) # returns a new function g(*args, *cotangents) # that computes vjps and (*args, cotangents) def get_vjp_fn_and_args_with_cotangents(f, sample, cotangents): args = tuple([sample.input] + list(sample.args)) kwargs = sample.kwargs flat_args, args_spec = tree_flatten(args) flat_cotangents, cotangents_spec = tree_flatten(cotangents) @functools.wraps(f) def wrapped(*args): assert len(args) == len(flat_args) + len(flat_cotangents) actual_args = args[:len(flat_args)] cotangents = args[len(flat_args):] actual_args = tree_unflatten(actual_args, args_spec) cotangents = tree_unflatten(cotangents, cotangents_spec) fn, primals = normalize_op_input_output3(f, actual_args, kwargs, flat_args, sample.output_process_fn_grad) _, vjp_fn = vjp(fn, *primals) return vjp_fn(cotangents) return wrapped, tuple(flat_args + flat_cotangents) # Returns a new function g(*args, *cotangents) that computes vjps and # sample (*args, *cotangents) def get_vjpfull_variant(f, sample): fn, primals = normalize_op_input_output(f, sample) result = fn(*primals) cotangents = _as_tuple( tree_map(lambda x: torch.randn_like(x, requires_grad=True), result)) num_primals = len(primals) args = (*primals, *cotangents) @functools.wraps(f) def wrapped(*args): primals = args[:num_primals] cotangents = args[num_primals:] result, vjp_fn = vjp(fn, *primals) if isinstance(result, torch.Tensor): assert len(cotangents) == 1 cotangents = cotangents[0] return vjp_fn(cotangents) return wrapped, args def get_jvp_variant(f, sample): # We want this higher-order variant of jvp, so that it can # be used to wrap vmap fn, primals = normalize_op_input_output(f, sample, requires_grad=False) tangents = _as_tuple( tree_map(lambda x: torch.randn_like(x), primals)) @functools.wraps(f) def wrapped(*args): tangents = args primals_out, tangents_out = jvp(fn, primals, tangents) if isinstance(primals_out, torch.Tensor): return (primals_out, tangents_out) else: flat_primals_out, _ = tree_flatten(primals_out) flat_tangents_out, _ = tree_flatten(tangents_out) return tuple(flat_primals_out + flat_tangents_out) return wrapped, tangents def get_jvp_variant_primals_tangents(f, sample): # We want this higher-order variant of jvp, so that it can # be used to wrap vmap fn, primals = normalize_op_input_output(f, sample, requires_grad=False) tangents = _as_tuple( tree_map(lambda x: torch.randn_like(x), primals)) @functools.wraps(f) def wrapped(*args): primals_in = args[:len(primals)] tangents_in = args[len(primals):] primals_out, tangents_out = jvp(fn, primals_in, tangents_in) if isinstance(primals_out, torch.Tensor): return (primals_out, tangents_out) else: flat_primals_out, _ = tree_flatten(primals_out) flat_tangents_out, _ = tree_flatten(tangents_out) return tuple(flat_primals_out + flat_tangents_out) return wrapped, primals + tangents def is_inplace(op, variant): if hasattr(variant, "__wrapped__"): return variant.__wrapped__ is op.get_inplace() return variant is op.get_inplace() vjp_fail = { xfail('tensor_split'), # data_ptr composite compliance } aliasing_ops = { 'T', 'broadcast_to', 'conj', 'contiguous', 'diagonal', # linalg.diagonal is an alias 'expand', 'flatten', 'imag', 'mH', # adjoint is an alias 'mT', 'movedim', # moveaxis is an alias 'narrow', 'permute', 'positive', # 'ravel', is composite implict autograd and may call clone 'real', 'reshape', 'resolve_conj', 'resolve_neg', 'select', 'squeeze', 'transpose', # swapdims and swapaxes are aliases 'unflatten', 'unfold', 'unsqueeze', 'view', 'view_as', 'view_as_complex', 'view_as_real', } aliasing_ops_list_return = { 'chunks', 'dsplit', 'hsplit', 'split', 'unbind', 'vsplit', # 'tensor_split' not composite compliant, see vjp_fail } class TestOperators(TestCase): @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_grad', vjp_fail.union({ xfail('linalg.eig'), # diagonal_scatter does not support complex xfail('chalf', '', device_type='cpu'), # RuntimeError: "sum_cpu" not implemented for 'ComplexHalf' skip('as_strided_scatter', ''), # silent incorrectness; seems flaky xfail('sparse.sampled_addmm', ''), # RuntimeError: Sparse CSR tensors do not have strides xfail('to_sparse', ''), # Could not run 'aten::sum.dim_IntList' })) @opsToleranceOverride('TestOperators', 'test_grad', ( tol1('nn.functional.binary_cross_entropy_with_logits', {torch.float32: tol(atol=1e-04, rtol=1e-04)}), )) def test_grad(self, device, dtype, op): if op.name in vjp_fail: self.skipTest("Skipped; Expected failures") return if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) if is_inplace(op, op.get_op()): self.skipTest("Skipped for redundancy. test_vjp handles in-place testing.") return for sample in samples: args = [sample.input] + list(sample.args) kwargs = sample.kwargs diff_argnums = tuple(i for i, arg in enumerate(args) if diff_arg(arg)) assert len(diff_argnums) > 0 diff_args = tuple(args[i] for i in diff_argnums) def wrapped_fn(*args, **kwargs): result = op(*args, **kwargs) if sample.output_process_fn_grad is not None: result = sample.output_process_fn_grad(result) # Reduce into single value for grad if isinstance(result, torch.Tensor): return result.sum() result = sum([res.sum() for res in result]) return result result = grad(wrapped_fn, diff_argnums)(*args, **kwargs) expected = _autograd_grad(_as_tuple(wrapped_fn(*args, **kwargs)), diff_args) self.assertEqual(result, expected) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_jvp', set({ # Composite ops that do bad things. Need to be fixed in PyTorch core. # RuntimeError: Cannot access data pointer of Tensor that doesn't have storage xfail('tensor_split'), # BUG: silent incorrectness: runs and produces numerical differences skip('nn.functional.max_unpool1d'), # fails everywhere except on mac skip('nn.functional.max_unpool2d'), # fails everywhere except on windows skip('nn.functional.max_unpool3d'), # fails everywhere except on mac xfail("native_batch_norm"), xfail('nn.functional.rrelu') # in-place test errors out with no formula implemented })) @opsToleranceOverride('TestOperators', 'test_jvp', ( tol1('nn.functional.conv_transpose3d', {torch.float32: tol(atol=1e-04, rtol=1.3e-06)}, device_type='cuda'), tol1('nn.functional.binary_cross_entropy_with_logits', {torch.float32: tol(atol=4e-04, rtol=4e-04)}), )) def test_jvp(self, device, dtype, op): # TODO: get rid of vjp_decomp when we add decomposition support to # PyTorch's forward-mode ad. Currently the decomposition support only # works for functorch.jvp VJP_DECOMP = { 'nn.functional.logsigmoid', } if op.name in VJP_DECOMP: fixme_ref_jvp_local = simulate_jvp else: fixme_ref_jvp_local = ref_jvp if not op.supports_forward_ad and op.name not in VJP_DECOMP: self.skipTest("Skipped! Forward AD not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) outplace_variant = op if not is_inplace(op, op.get_op()) else None inplace_variant = op.inplace_variant if op.supports_inplace_autograd else None for sample in samples: args = (sample.input,) + sample.args kwargs = sample.kwargs if outplace_variant: self.jvp_opinfo_test(outplace_variant, args, kwargs, sample.output_process_fn_grad, clone_inputs=False, fixme_ref_jvp_local=fixme_ref_jvp_local) if is_valid_inplace_sample_input(sample, op, inplace_variant): self.jvp_opinfo_test(inplace_variant, args, kwargs, sample.output_process_fn_grad, clone_inputs=True, fixme_ref_jvp_local=fixme_ref_jvp_local) def jvp_opinfo_test(self, fn, args, kwargs, output_process_fn, clone_inputs, fixme_ref_jvp_local): # NB: we used requires_grad=True to determine where the primals are, # but don't need that information otherwise fn, primals = normalize_op_input_output2( fn, args, kwargs, output_process_fn, requires_grad=True) orig_primals = tree_map(lambda x: x.detach(), primals) orig_tangents = tree_map(lambda x: torch.randn_like(x), primals) def maybe_clone_inputs(): if clone_inputs: primals = tree_map(torch.clone, orig_primals) tangents = tree_map(torch.clone, orig_tangents) return primals, tangents return orig_primals, orig_tangents primals, tangents = maybe_clone_inputs() expected_primal_outs, expected_tangent_outs = \ fixme_ref_jvp_local(fn, primals, tangents) primals, tangents = maybe_clone_inputs() primal_outs, tangent_outs = jvp(fn, primals, tangents) self.assertEqual(primal_outs, expected_primal_outs) self.assertEqual(tangent_outs, expected_tangent_outs) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_vjp', vjp_fail.union({ skip('as_strided_scatter', ''), # silent incorrectness; also might be flaky xfail('sparse.sampled_addmm', ''), })) @opsToleranceOverride('TestOperators', 'test_vjp', ( tol1('nn.functional.conv_transpose3d', {torch.float32: tol(atol=5e-05, rtol=9e-05)}, device_type='cuda'), tol1('nn.functional.binary_cross_entropy_with_logits', {torch.float32: tol(atol=1e-04, rtol=1e-04)}), )) def test_vjp(self, device, dtype, op): if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) def _test(_op, inplace=False): for sample in samples: if inplace and not is_valid_inplace_sample_input(sample, op, op.inplace_variant): continue fn, primals = normalize_op_input_output(_op, sample) result = fn(*primals) cotangents = tree_map(lambda x: torch.randn_like(x), result) out, vjp_fn = vjp(fn, *primals) self.assertEqual(out, result) result_vjps = vjp_fn(cotangents) _, vjp_fn = ref_vjp(fn, *primals) expected_vjps = vjp_fn(cotangents) self.assertEqual(result_vjps, expected_vjps) _test(op) for a_op in op.aliases: _test(a_op) if op.inplace_variant: def f(inp, *args, **kwargs): return op.inplace_variant(inp.clone(), *args, **kwargs) _test(f, inplace=True) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_vjpvjp', vjp_fail.union({ skip('nn.functional.max_unpool1d'), # silent incorrectness; Flaky skip('nn.functional.max_unpool2d'), # silent incorrectness; Flaky xfail('nn.functional.ctc_loss'), # Not Implemented xfail('native_layer_norm', ''), # Expected a proper Tensor but got None for argument #1 'other' xfail('sparse.sampled_addmm', ''), # sparse tensors have no strides # AssertionError: Tensor-likes are not close! # Mismatched elements: 1 / 15 (6.7%) # Greatest absolute difference: 24.0 at index (2, 4) (up to 1e-05 allowed) # Greatest relative difference: 1.7933241714393998e-06 at index (2, 4) (up to 1.3e-06 allowed) # The failure occurred for item [0] xfail('masked.prod') })) @opsToleranceOverride('TestOperators', 'test_vjpvjp', ( tol1('nn.functional.conv_transpose3d', {torch.float32: tol(atol=5e-05, rtol=9e-05)}, device_type='cuda'), tol1('prod', {torch.float32: tol(atol=2e-05, rtol=1e-04)}), tol1('masked.cumprod', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('cumprod', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('linalg.vander', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol2('linalg.det', 'singular', {torch.float32: tol(atol=2e-05, rtol=2e-05)}), )) def test_vjpvjp(self, device, dtype, op): if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return if not op.supports_gradgrad: self.skipTest("Skipped! Operation does not support gradgrad") return samples = op.sample_inputs(device, dtype, requires_grad=True) def test(_op, inplace=False): for sample in samples: if inplace and not is_valid_inplace_sample_input(sample, op, op.inplace_variant): continue fn, args = get_vjpfull_variant(_op, sample) result = fn(*args) cotangents = tree_map(lambda x: torch.randn_like(x), result) # Compute vjp of vjp _, vjp_fn = vjp(fn, *args) result_vjps = vjp_fn(cotangents) # Compute ref_vjp of vjp. We could have done ref_vjp of ref_vjp, # but since we're confident that vjp works by itself, this is # an equivalent way to test that. _, vjp_fn = ref_vjp(fn, *args) expected_vjps = vjp_fn(cotangents) self.assertEqual(result_vjps, expected_vjps) test(op) if op.inplace_variant: def fn(inp, *args, **kwargs): return op.inplace_variant(inp.clone(), *args, **kwargs) test(fn, inplace=True) @skipOps('TestOperators', 'test_vmapvjpvjp', vjp_fail.union({ skip("atleast_1d"), # Takes too long skip("atleast_2d"), # Takes too long skip("atleast_3d"), # Takes too long xfail("as_strided"), # incorrect output xfail("as_strided_scatter"), # incorrect output skip("bernoulli"), # calls random op xfail("bfloat16"), # rank 4 tensor for channels_last xfail("chalf"), # rank 4 tensor for channels_last xfail("double"), # rank 4 tensor for channels_last xfail("float"), # rank 4 tensor for channels_last xfail("half"), # rank 4 tensor for channels_last # It looks like you're either (1) calling .item() on a Tensor or # (2) attempting to use a Tensor in some data-dependent control flow or # (3) encountering this error in PyTorch internals. xfail("index_reduce"), xfail("linalg.eig"), # vmap over torch.allclose xfail("linalg.eigvals"), # vmap over torch.allclose xfail("linalg.householder_product"), # vmap: inplace into a regular tensor xfail("nanquantile", device_type='cpu'), # vmap not implemented for at::equal. xfail("native_layer_norm"), # vmap: inplace into a regular tensor # got a batched tensor as input while the running_mean or running_var, # which will be updated in place, were not batched. xfail("nn.functional.batch_norm"), xfail("nn.functional.binary_cross_entropy"), # vmap: inplace into a regular tensor xfail("nn.functional.ctc_loss"), # derivate not implemented for _ctc_loss_backward skip("nn.functional.dropout"), # calls random op skip("nn.functional.dropout2d"), # calls random op skip("nn.functional.dropout3d"), # calls random op skip("nn.functional.feature_alpha_dropout", "with_train"), # calls random op skip("nn.functional.fractional_max_pool2d"), # calls random op skip("nn.functional.fractional_max_pool3d"), # calls random op skip('nn.functional._scaled_dot_product_attention'), # randomness # It looks like you're either (1) calling .item() on a Tensor or # (2) attempting to use a Tensor in some data-dependent control flow or # (3) encountering this error in PyTorch internals. xfail("nn.functional.gaussian_nll_loss"), # got a batched tensor as input while the running_mean or running_var, # which will be updated in place, were not batched. xfail("nn.functional.instance_norm"), xfail("nn.functional.layer_norm"), # vmap: inplace into a regular tensor # RuntimeError: NYI: querying is_contiguous inside of vmap # for memory_format other than torch.contiguous_formats xfail("nn.functional.max_pool2d"), # RuntimeError: NYI: Tensor.clone(memory_format) inside vmap is only # supported with memory_format torch.preserve_format or # torch.contiguous_format (got ChannelsLast) xfail("nn.functional.max_unpool2d"), # RuntimeError: NYI: Tensor.clone(memory_format) inside vmap is only # supported with memory_format torch.preserve_format # or torch.contiguous_format (got ChannelsLast)s xfail("nn.functional.max_unpool2d", "grad"), xfail("nn.functional.rrelu"), # RuntimeError: vmap: we do not yet support aten::rrelu_with_noise. xfail("normal"), # calls random op xfail("normal", "number_mean"), # calls random op xfail("pca_lowrank"), # calls random op xfail("put"), # vmap: inplace into a regular tensor xfail("quantile", device_type='cpu'), # Batching rule not implemented for `at::equal` xfail("scatter_reduce", "prod"), # vmap (looks like you are calling item/data-dependent) xfail("sparse.sampled_addmm"), # RuntimeError: Sparse CSR tensors do not have strides xfail("svd_lowrank"), # calls random op xfail("take"), # vmap: inplace into a regular tensor xfail("to"), # rank 4 tensor for channels_last xfail("view_as_complex"), # RuntimeError: Tensor must have a last dimension with stride 1 xfail("masked.softmax", device_type='cuda'), # Mismatch in values! xfail("masked.softmin", device_type='cuda'), # Mismatch in values! # got a batched tensor as input while the running_mean or running_var, # which will be updated in place, were not batched. xfail("nn.functional.batch_norm", 'without_cudnn'), # view doesn't work on sparse xfail("to_sparse"), xfail("native_batch_norm"), })) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) @opsToleranceOverride('TestOperators', 'test_vmapvjpvjp', ( tol1('linalg.svd', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('linalg.lu_factor', {torch.float32: tol(atol=2e-03, rtol=2e-02)}), tol1('svd', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), )) def test_vmapvjpvjp(self, device, dtype, op): # Since, we test `vjpvjp` independently, # for this test, we just verify that vmap # of `vjpvjp` is correct. if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return if not op.supports_gradgrad: self.skipTest("Skipped! Operation does not support gradgrad") return samples = op.sample_inputs(device, dtype, requires_grad=True) # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return for sample in samples: fn, args = get_vjpfull_variant(op, sample) result = fn(*args) cotangents = tree_map(lambda x: torch.randn_like(x), result) cotangents, _ = tree_flatten(cotangents) num_args = len(args) args_and_cotangents = tuple(args) + tuple(cotangents) def vjp_of_vjp(*args_and_cotangents): args = args_and_cotangents[:num_args] cotangents = args_and_cotangents[num_args:] result, vjp_fn = vjp(fn, *args) result_vjps = vjp_fn(cotangents) result, _ = tree_flatten(result) result_vjps, _ = tree_flatten(result_vjps) return (*result, *result_vjps) is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs) generator = get_fallback_and_vmap_exhaustive( vjp_of_vjp, args_and_cotangents, {}, is_batch_norm_and_training=is_batch_norm_and_training) for loop_out, batched_out in generator: self.assertEqual(loop_out, batched_out) vmapvjp_fail = vjp_fail.union({ # -------------------- ALLOWED FAILURES -------------------------------- # The following are not bugs and are expected behavior xfail('masked_select'), # Not possible due to dynamic shapes skip('bernoulli'), # randomness skip('normal', ''), # randomness skip('normal', 'number_mean'), # randomness skip('nn.functional.rrelu'), # randomness skip('nn.functional.feature_alpha_dropout', 'with_train'), # randomness skip('nn.functional.feature_alpha_dropout', 'without_train'), # randomness skip('nn.functional.dropout'), # randomness skip('nn.functional.dropout2d'), # randomness skip('nn.functional.dropout3d', ''), # randomness skip('nn.functional._scaled_dot_product_attention'), # randomness xfail('as_strided'), # as_strided is too wild for us to support, wontfix xfail('index_put', ''), # not possible due to dynamic shapes; we support a subset xfail('masked_scatter'), # dynamic xfail('nn.functional.fractional_max_pool2d'), # random xfail('nn.functional.fractional_max_pool3d'), # random xfail('take'), # dynamic xfail('pca_lowrank', ''), # randomness xfail('svd_lowrank', ''), # randomness xfail('to_sparse', ''), # non-dense output skip('to'), # RuntimeError: required rank 4 tensor to use channels_last format # ---------------------------------------------------------------------- # ---------------------------- BUGS ------------------------------------ # All of the following are bugs and need to be fixed skip('linalg.svdvals'), # # really annoying thing where it passes correctness check but not has_batch_rule skip("native_batch_norm"), xfail('__getitem__', ''), # dynamic error xfail('linalg.eig'), # Uses aten::allclose xfail('linalg.householder_product'), # needs select_scatter xfail('nanquantile', device_type='cpu'), # checks q via a .item() call xfail('nn.functional.gaussian_nll_loss'), # checks var for if any value < 0 xfail('narrow'), # .item() call xfail('quantile', device_type='cpu'), # checks q via a .item() call xfail('view_as_complex'), # Tensor must have a last dimension with stride 1 # required rank 4 tensor to use channels_last format xfail('bfloat16'), xfail('double'), xfail('float'), xfail('half'), xfail('chalf', ''), xfail('scatter_reduce', 'prod'), # item call # Batching rule not implemented for aten::_use_cudnn_ctc_loss.Tensor xfail('nn.functional.ctc_loss', device_type='cuda'), # NYI: querying is_contiguous inside of vmap for memory_format other than torch.contiguous_format xfail('nn.functional.max_unpool2d'), xfail('nn.functional.max_unpool2d', 'grad'), xfail('sparse.sampled_addmm', ''), xfail('as_strided_scatter', ''), # calls as_strided xfail('index_reduce', ''), # .item() call # --------------------------------------------------------------------- }) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) @opsToleranceOverride('TestOperators', 'test_vmapvjp', ( tol1('linalg.svd', {torch.float32: tol(atol=1.5e-04, rtol=1e-04)}, device_type="cuda"), tol1('svd', {torch.float32: tol(atol=1.5e-04, rtol=1e-04)}, device_type="cuda"), )) @skipOps('TestOperators', 'test_vmapvjp', vmapvjp_fail) def test_vmapvjp(self, device, dtype, op): if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return for sample in samples: cotangents = get_sample_cotangents(op, sample) fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents) is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs) generator = get_fallback_and_vmap_exhaustive( fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training) for loop_out, batched_out in generator: self.assertEqual(loop_out, batched_out) vmapjvpall_fail = { # -------------------- ALLOWED FAILURES -------------------------------- # The following are expected (not a bug) skip('bernoulli', ''), # randomness skip('nn.functional.dropout'), # randomness skip('nn.functional.rrelu'), # randomness skip('nn.functional.dropout2d', ''), skip('nn.functional.dropout3d', ''), skip('nn.functional._scaled_dot_product_attention'), # randomness skip('nn.functional.feature_alpha_dropout', 'without_train'), skip('nn.functional.feature_alpha_dropout', 'with_train'), xfail('nn.functional.fractional_max_pool2d'), # Cannot access data pointer of Tensor that doesn't have storage xfail('nn.functional.fractional_max_pool3d'), # Cannot access data pointer of Tensor that doesn't have storage # Not actually a problem: embedding with max_norm mutates the weight # and causes different runs to produce different results. # skip because this is flaky depending on what the max_norm is! skip('nn.functional.embedding', ''), skip('to'), # RuntimeError: required rank 4 tensor to use channels_last format # ---------------------------------------------------------------------- # ---------------------------- BUGS ------------------------------------ # The following are bugs that we should fix decorate('nn.functional.conv2d', decorator=unittest.skipIf(IS_ARM64, "Fails on M1")), skip('nn.functional.max_pool1d'), # fails on cpu, runs on cuda xfail('masked.mean'), # silent incorrectness (nan difference) xfail('nn.functional.soft_margin_loss', ''), # soft_margin_loss_backward does not support forward-ad xfail('tensor_split'), # data_ptr composite compliance xfail('quantile'), # at::equal batching rule (cpu), also, in-place vmap (cuda) skip('as_strided'), # Test runner cannot handle this xfail('nn.functional.gaussian_nll_loss'), # .item or data-dependent control flow xfail('scatter'), # forward-mode AD does not support at::scatter xfail('nanquantile'), # at::equal batching rule (cpu), also, in-place vmap (cuda) xfail('view_as_complex'), # Tensor must have a last dimension with stride 1 skip('pca_lowrank', ''), # randomness skip('svd_lowrank', ''), # randomness xfail('double'), # required rank 4 tensor to use channels_last format # potential silent incorrectness skip('nn.functional.max_unpool1d'), # Flaky, seems to sometimes his max_unpool2d skip('nn.functional.max_unpool2d'), # fails everywhere except on mac skip('nn.functional.max_unpool3d'), # fails everywhere except on mac # erroring because running_mean and running_var aren't differentiable xfail('nn.functional.batch_norm'), xfail('nn.functional.batch_norm', 'without_cudnn'), xfail("native_batch_norm"), # ---------------------------------------------------------------------- } @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @opsToleranceOverride('TestOperators', 'test_vmapjvpall', ( tol1('nn.functional.conv_transpose3d', {torch.float32: tol(atol=2e-04, rtol=9e-3)}, device_type='cuda'), tol1('linalg.householder_product', {torch.float32: tol(atol=2e-04, rtol=9e-3)}, device_type='cuda'), tol1('linalg.householder_product', {torch.float32: tol(atol=2e-04, rtol=1e-4)}, device_type='cpu'), )) @skipOps('TestOperators', 'test_vmapjvpall', vmapjvpall_fail) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) # This is technically a superset of test_vmapjvp. We should either delete test_vmapjvp # or figure out if we can split vmapjvpall. It's useful to keep test_vmapjvp intact # because that coresponds to "batched forward-mode AD" testing in PyTorch core def test_vmapjvpall(self, device, dtype, op): if is_inplace(op, op.get_op()): # TODO: test in-place self.skipTest("Skipped! NYI: inplace-testing not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=False) if not op.supports_forward_ad: self.skipTest("Skipped! Forward AD not supported.") return for sample in samples: arg_values = [sample.input] + list(sample.args) kwarg_values = sample.kwargs args = tuple(arg_values) + tuple(kwarg_values) fn, args = get_jvp_variant_primals_tangents(op, sample) is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values) generator = get_fallback_and_vmap_exhaustive( fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training) for loop_out, batched_out in generator: self.assertEqual(loop_out, batched_out) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_vmapjvpall_has_batch_rule', vmapjvpall_fail.union({ skip('to'), # RuntimeError: required rank 4 tensor to use channels_last format xfail('nn.functional.huber_loss'), xfail('lu'), xfail('cumprod'), xfail('masked_fill'), xfail('copysign'), xfail('complex'), skip('masked.mean'), # ??? xfail('masked_scatter'), xfail('index_fill'), xfail('put'), xfail('take'), xfail('nn.functional.max_pool3d'), xfail('vdot'), xfail('nanmean'), xfail('nansum'), xfail('nn.functional.feature_alpha_dropout', 'without_train'), xfail('linalg.lu_factor', ''), xfail('nn.functional.dropout2d', ''), xfail('pca_lowrank', ''), xfail('svd_lowrank', ''), xfail('linalg.lu_factor_ex', ''), xfail('nn.functional.feature_alpha_dropout', 'with_train'), xfail('special.log_ndtr', ''), xfail('fft.ihfft2'), # conj_physical fallback xfail('fft.ihfftn'), # conj_physical fallback xfail('polar'), # complex fallback xfail('nn.functional.max_unpool3d', 'grad'), xfail('nn.functional.smooth_l1_loss', ''), xfail('nn.functional.max_unpool2d', 'grad'), xfail('nn.functional.soft_margin_loss', ''), xfail('nn.functional.max_unpool1d', 'grad'), xfail('nn.functional.embedding', ''), xfail('scatter_reduce', "sum"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "mean"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "amin"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "amax"), # aten::scatter_reduce.two hit the vmap fallback xfail('lu_unpack'), xfail('nn.functional.glu'), xfail('nn.functional.bilinear'), # trilinear doesn't have batching rule xfail('linalg.lu', ''), xfail('linalg.lu_solve', ''), xfail('nn.functional.dropout3d', ''), xfail('as_strided_scatter', ''), xfail('masked.cumprod', ''), xfail('linalg.vecdot', ''), })) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) def test_vmapjvpall_has_batch_rule(self, device, dtype, op): if is_inplace(op, op.get_op()): # TODO: test in-place self.skipTest("Skipped! NYI: inplace-testing not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=False) if not op.supports_forward_ad: self.skipTest("Skipped! Forward AD not supported.") return def test(): for sample in samples: arg_values = [sample.input] + list(sample.args) kwarg_values = sample.kwargs args = tuple(arg_values) + tuple(kwarg_values) fn, args = get_jvp_variant_primals_tangents(op, sample) is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values) for loop_out, batched_out in get_fallback_and_vmap_exhaustive( fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False): pass check_vmap_fallback(self, test, op, dry_run=False) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) @skipOps('TestOperators', 'test_vmapvjp_has_batch_rule', vmapvjp_fail.union({ skip('to'), # RuntimeError: required rank 4 tensor to use channels_last format xfail('view_as_complex'), xfail('complex'), xfail('copysign'), xfail('cummax'), xfail('cummin'), xfail('cumprod'), xfail('nansum'), xfail('nanmean'), xfail('narrow'), # Batching rule not implemented for `narrow.Tensor` (and view op) xfail('special.log_ndtr'), xfail('index_copy'), xfail('index_fill'), xfail('linalg.eig'), xfail('linalg.householder_product'), xfail('lu'), xfail('lu_solve'), xfail('lu_unpack'), xfail('masked_fill'), xfail('masked_scatter'), xfail('masked_select'), xfail('nanquantile'), xfail('put'), xfail('scatter_reduce', "sum"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "mean"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "amin"), # aten::scatter_reduce.two hit the vmap fallback xfail('scatter_reduce', "amax"), # aten::scatter_reduce.two hit the vmap fallback xfail('quantile'), xfail('renorm'), xfail('take'), xfail('tensor_split'), xfail('to_sparse'), xfail('unfold'), xfail('vdot'), xfail('nn.functional.dropout'), xfail('fft.ihfft2'), xfail('fft.ihfftn'), xfail('nn.functional.gaussian_nll_loss'), xfail('nn.functional.huber_loss'), xfail('nn.functional.bilinear'), xfail('nn.functional.fractional_max_pool3d'), xfail('nn.functional.ctc_loss'), xfail('as_strided'), xfail('stft'), xfail('nn.functional.rrelu'), xfail('nn.functional.embedding_bag'), xfail('nn.functional.max_pool3d'), xfail('nn.functional.fractional_max_pool2d'), xfail('linalg.lu_factor', ''), xfail('nn.functional.feature_alpha_dropout', 'with_train'), xfail('pca_lowrank', ''), xfail('nn.functional.dropout2d', ''), xfail('nn.functional.feature_alpha_dropout', 'without_train'), xfail('svd_lowrank', ''), xfail('linalg.lu_factor_ex', ''), xfail('nn.functional.max_unpool2d', ''), xfail('nn.functional.multi_margin_loss', ''), xfail('nn.functional.multilabel_margin_loss', ''), xfail('nn.functional.pdist', ''), xfail('nn.functional.smooth_l1_loss', ''), xfail('scatter_reduce', 'prod'), xfail('nn.functional.max_unpool1d', ''), xfail('nn.functional.max_unpool3d', ''), xfail('nn.functional.max_unpool3d', 'grad'), xfail('nn.functional.soft_margin_loss', ''), xfail('nn.functional.max_unpool1d', 'grad'), xfail('nn.functional.max_unpool2d', 'grad'), xfail('linalg.lu', ''), xfail('linalg.lu_solve', ''), xfail('chalf', ''), xfail('index_reduce', ''), xfail('linalg.vander', ''), xfail('nn.functional.dropout3d', ''), xfail('as_strided_scatter', ''), xfail('segment_reduce', 'offsets'), xfail('masked.cumprod', ''), xfail('linalg.vecdot', ''), xfail('segment_reduce', 'lengths'), xfail('sparse.sampled_addmm', ''), xfail("native_batch_norm"), })) def test_vmapvjp_has_batch_rule(self, device, dtype, op): if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return def test(): for sample in samples: cotangents = get_sample_cotangents(op, sample) fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents) is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs) for loop_out, batched_out in get_fallback_and_vmap_exhaustive( fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False): pass for a_op in op.aliases: fn, args = get_vjp_fn_and_args_with_cotangents(a_op, sample, cotangents) for loop_out, batched_out in get_fallback_and_vmap_exhaustive( fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False): pass check_vmap_fallback(self, test, op, dry_run=False) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_vjpvmap', vjp_fail.union({ skip('bernoulli', ''), # vjpvmap testing can't handle randomness skip('normal', ''), # vjpvmap testing can't handle randomness skip('normal', 'number_mean'), # vjpvmap testing can't handle randomness skip('nn.functional.rrelu'), # randomness skip('nn.functional.feature_alpha_dropout', 'with_train'), # randomness skip('nn.functional.feature_alpha_dropout', 'without_train'), # randomness skip('to'), # RuntimeError: required rank 4 tensor to use channels_last format skip('to_sparse', ''), # non-dense output # fallback path doesn't work # All of the following are bugs and need to be fixed xfail('__getitem__', ''), xfail('index_put', ''), xfail('view_as_complex'), xfail('nn.functional.gaussian_nll_loss'), xfail('masked_select'), xfail('narrow'), # Batching rule not implemented for `narrow.Tensor` (and view op) skip('nn.functional.fractional_max_pool3d'), # generator works on cpu, fails on cuda xfail('__rpow__'), # https://github.com/pytorch/functorch/issues/617 skip('nn.functional.fractional_max_pool2d'), # generator works on cpu, fails on cuda xfail('column_stack', ''), xfail('nn.functional.dropout2d', ''), xfail('svd_lowrank', ''), xfail('pca_lowrank', ''), xfail('clamp'), xfail('cross'), # The defaults of this op are *very* weird. No wonder it doesn't work # something weird happening with channels_last xfail('bfloat16'), xfail('double'), xfail('float'), xfail('half'), xfail('nn.functional.dropout3d', ''), xfail('as_strided_scatter', ''), xfail('sparse.sampled_addmm', ''), xfail("native_batch_norm"), })) def test_vjpvmap(self, device, dtype, op): # NB: there is no vjpvmap_has_batch_rule test because that is almost # certainly redundant with the vmap_has_batch_rule test in test_vmap.py # one-off skip if op.name == 'nn.functional.dropout': self.skipTest("Skipped!") if not op.supports_autograd: # If the op doesn't support autograd, vmap(op) won't either self.skipTest("Skipped! Autograd not supported.") return # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) batch_norm_fns = ("nn.functional.batch_norm", "nn.functional.instance_norm") # instance norm calls batch norm is_batch_norm = op.name in batch_norm_fns for sample in samples: args = [sample.input] + list(sample.args) kwargs = sample.kwargs is_batch_norm_and_training = is_batch_norm and is_batch_norm_training(op.name, kwargs) generator = generate_vmap_inputs(args, kwargs, is_batch_norm_and_training=is_batch_norm_and_training) for batched_args, in_dims, kwargs in generator: vmapped_op = vmap(op, in_dims) fn, primals = normalize_op_input_output2(vmapped_op, batched_args, kwargs, sample.output_process_fn_grad) result = fn(*primals) cotangents = tree_map(lambda x: torch.randn_like(x), result) _, vjp_fn = vjp(fn, *primals) result_vjps = vjp_fn(cotangents) _, vjp_fn = ref_vjp(fn, *primals) expected_vjps = vjp_fn(cotangents) self.assertEqual(result_vjps, expected_vjps) def _compare_jacobians_of_vjp(self, fn, cotangents_and_primals, argnums=None, atol_rtol=None): if argnums is None: argnums = tuple(range(len(cotangents_and_primals))) def get_vjp(cotangents, *primals): _, vjp_fn = vjp(fn, *primals) return vjp_fn(cotangents) jacobian_jvp = jacfwd(get_vjp, argnums)(*cotangents_and_primals) jacobian_vjp = jacrev(get_vjp, argnums)(*cotangents_and_primals) # For dtype changing operations, the jacobians have different dtype. jacobian_jvp = tree_map(lambda x: x.to(torch.float), jacobian_jvp) jacobian_vjp = tree_map(lambda x: x.to(torch.float), jacobian_vjp) if atol_rtol is not None: (atol, rtol) = atol_rtol self.assertEqual(jacobian_jvp, jacobian_vjp, atol=atol, rtol=rtol) else: self.assertEqual(jacobian_jvp, jacobian_vjp) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @skipOps('TestOperators', 'test_jvpvjp', vjp_fail.union({ xfail('to_sparse', ''), # NYI # RuntimeError: Trying to set a forward gradient that has a different size than that of the original Tensor, # this is not supported. Tensor is of size [5, 2, 3] while the given forward gradient is of size [1, 2, 3]. xfail('normal', ''), xfail('cdist', ''), # NYI: forward-AD for _cdist_forward xfail('cholesky', ''), # NYI: forward-AD for cholesky xfail('logcumsumexp', ''), # NYI: forward-AD for logcumsumexp xfail('nn.functional.embedding_bag', ''), # NYI: forward-AD for _embedding_bag xfail('nn.functional.grid_sample', ''), # NYI: forward AD for grid_sampler_2d xfail('nn.functional.hardsigmoid', ''), # NYI: forward AD for hardsigmoid_backward xfail('nn.functional.huber_loss', ''), # NYI: forward AD for huber_loss_backward xfail('nn.functional.logsigmoid', ''), # not differentiable w.r.t. buffer xfail('renorm', ''), # NYI: forward AD for renorm xfail('symeig', ''), # NYI: forward AD for symeig xfail('nn.functional.multilabel_margin_loss', ''), # NYI: multilabel_margin_loss_forward xfail('nn.functional.multilabel_soft_margin_loss', ''), # NYI: log_sigmoid_backward xfail('nn.functional.soft_margin_loss', ''), # NYI: forward-AD for log_sigmoid_backward xfail('nn.functional.ctc_loss', ''), # NYI: forward-AD for _ctc_loss xfail('nn.functional.pdist', ''), # NYI: forward-AD with _pdist_forward xfail('nn.functional.multi_margin_loss', ''), # NYI: forward AD with multi_margin_loss skip('linalg.householder_product', '', device_type='cuda'), # flaky, I'm not sure why xfail('sparse.sampled_addmm', ''), # Sparse tensors have no strides skip('as_strided_scatter', ''), # seems flaky xfail('segment_reduce', 'offsets'), # NYI: forward-AD for segment_reduce xfail('index_reduce', ''), # NYI: forward-AD for index_reduce xfail('segment_reduce', 'lengths'), # NYI: forward-AD for segment_reduce })) @opsToleranceOverride('TestOperators', 'test_jvpvjp', ( tol1('masked.prod', {torch.float32: tol(atol=1e-04, rtol=1.3e-05)}), tol1('masked.cumprod', {torch.float32: tol(atol=1e-04, rtol=1e-04)}), tol1('cumprod', {torch.float32: tol(atol=1e-04, rtol=1.3e-05)}, device_type='cuda'), tol1('linalg.vander', {torch.float32: tol(atol=1e-04, rtol=1.3e-05)}, device_type='cuda'), )) def test_jvpvjp(self, device, dtype, op): if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return for sample in samples: fn, primals = normalize_op_input_output(op, sample) result = fn(*primals) cotangents = tree_map(lambda x: torch.randn_like(x), result) primals_tangents = tree_map(lambda x: torch.randn_like(x), primals) cotangents_tangents = tree_map(lambda x: torch.randn_like(x), cotangents) def push_vjp(primals, cotangents): _, vjp_fn = vjp(fn, *primals) return vjp_fn(cotangents) result = jvp(push_vjp, (primals, cotangents), (primals_tangents, cotangents_tangents)) self.assertEqual(len(result), 2) def tree_map2(fn, first, second): flat_first, spec_first = tree_flatten(first) flat_second, spec_second = tree_flatten(second) assert spec_first == spec_second flat_result = [fn(f, s) for f, s in zip(flat_first, flat_second)] return tree_unflatten(flat_result, spec_first) def reference(primals, cotangents, primals_tangents, cotangents_tangents): with fwAD.dual_level(): primal_duals = tree_map2(fwAD.make_dual, primals, primals_tangents) _, vjp_fn = ref_vjp(fn, *primal_duals) cotangent_duals = tree_map2(fwAD.make_dual, cotangents, cotangents_tangents) result = vjp_fn(cotangent_duals) flat_result, spec = tree_flatten(result) primals_out, tangents_out = zip(*[fwAD.unpack_dual(r) for r in flat_result]) tangents_out = [t if t is not None else torch.zeros_like(p) for p, t in zip(primals_out, tangents_out)] expected = (tree_unflatten(primals_out, spec), tree_unflatten(tangents_out, spec)) return expected expected = reference(primals, cotangents, primals_tangents, cotangents_tangents) self.assertEqual(result, expected) @skipOps('TestOperators', 'test_vmapjvpvjp', vjp_fail.union({ # Following operatos take too long, hence skipped skip('atleast_1d'), skip('atleast_2d'), skip('atleast_3d'), skip('meshgrid', 'list_of_tensors'), skip('meshgrid', 'variadic_tensors'), skip('broadcast_tensors'), skip('linalg.lstsq'), skip('nn.functional.bilinear'), skip('native_layer_norm'), # Potential bugs/errors xfail('as_strided'), # AssertionError: Tensor-likes are not close! xfail('as_strided_scatter'), # AssertionError: Tensor-likes are not close! xfail('bernoulli'), # calls random op xfail('bfloat16'), # required rank 4 tensor to use channels_last format xfail('cdist'), # Forward AD not implemented and no decomposition xfail('chalf'), # required rank 4 tensor to use channels_last format xfail('cholesky'), # Forward AD not implemented and no decomposition xfail('double'), # required rank 4 tensor to use channels_last format xfail('float'), # required rank 4 tensor to use channels_last format xfail('half'), # required rank 4 tensor to use channels_last format xfail('index_reduce'), # Forward AD not implemented and no decomposition xfail('linalg.eig'), # vmap over torch.allclose isn't supported yet. # AssertionError: Tensor-likes are not close! # Mismatched elements: 2 / 120 (1.7%) # Greatest absolute difference: 0.09438323974609375 # Greatest relative difference: 0.00115722746596277 xfail('linalg.householder_product', device_type='cuda'), xfail('logcumsumexp'), # Forward AD not implemented and no decomposition xfail('mvlgamma', 'mvlgamma_p_1'), # vmap: inplace into a regular tensor xfail('mvlgamma', 'mvlgamma_p_3'), # vmap: inplace into a regular tensor xfail('mvlgamma', 'mvlgamma_p_5'), # vmap: inplace into a regular tensor xfail('nanquantile'), # Batching rule not implemented for aten::equal # RuntimeError: Batch norm got a batched tensor as input while the # running_mean or running_var, which will be updated in place, # were not batched. xfail('nn.functional.batch_norm'), xfail('nn.functional.batch_norm', 'without_cudnn'), xfail('nn.functional.binary_cross_entropy'), # vmap: inplace into a regular tensor xfail("nn.functional.ctc_loss"), # ForwardAD not implemented and no decomposition xfail('nn.functional.dropout2d'), # calls random op xfail('nn.functional.dropout3d'), # calls random op xfail('nn.functional.dropout'), # calls random op skip('nn.functional._scaled_dot_product_attention'), # randomness xfail('nn.functional.embedding_bag'), # Forward AD not implemented and no decomposition xfail('nn.functional.feature_alpha_dropout', 'with_train'), # calls random op xfail('nn.functional.fractional_max_pool2d'), # calls random op xfail('nn.functional.fractional_max_pool3d'), # calls random op xfail('nn.functional.gaussian_nll_loss'), # data depenedant flow xfail('nn.functional.grid_sample'), # Forward AD not implemented and no decomposition xfail('nn.functional.hardsigmoid'), # Forward AD not implemented and no decomposition xfail('nn.functional.hinge_embedding_loss'), # vmap: inplace into a regular tensor xfail('nn.functional.huber_loss'), # Forward AD not implemented and no decomposition # RuntimeError: Batch norm got a batched tensor as input while the # running_mean or running_var, which will be updated in place, # were not batched. xfail('nn.functional.instance_norm'), xfail('nn.functional.logsigmoid'), # Forward AD not implemented and no decomposition # NYI: Tensor.clone(memory_format) inside vmap is only supported with # memory_format torch.preserve_format or torch.contiguous_format (got ChannelsLast) xfail('nn.functional.max_unpool2d'), xfail('nn.functional.max_unpool2d', 'grad'), xfail('nn.functional.multi_margin_loss'), # Forward AD not implemented and no decomposition xfail('nn.functional.multilabel_margin_loss'), # Forward AD not implemented and no decomposition xfail('nn.functional.multilabel_soft_margin_loss'), # Forward AD not implemented and no decomposition xfail('nn.functional.pdist'), # Forward AD not implemented and no decomposition xfail('nn.functional.rrelu'), # vmap: we do not yet support aten::rrelu_with_noise. xfail('nn.functional.soft_margin_loss'), # Forward AD not implemented and no decomposition xfail('normal'), # calls random op xfail('normal', 'number_mean'), # calls random op xfail('pca_lowrank'), # calls random op xfail('quantile'), # Batching rule not implemented for aten::equal xfail('renorm'), # Forward AD not implemented and no decomposition xfail('scatter_reduce', 'prod'), # Forward AD not implemented and no decomposition xfail('segment_reduce', 'lengths'), # Forward AD not implemented and no decomposition xfail('segment_reduce', 'offsets'), # Forward AD not implemented and no decomposition xfail('sparse.sampled_addmm'), # RuntimeError: Sparse CSR tensors do not have strides xfail('svd_lowrank'), # calls random op xfail('symeig'), # Forward AD not implemented and no decomposition xfail('take'), # vmap: inplace into regular tensor xfail('to'), # RuntimeError: required rank 4 tensor to use channels_last format xfail('to_sparse'), # Forward AD not implemented and no decomposition xfail('view_as_complex'), # RuntimeError: Tensor must have a last dimension with stride 1 # RuntimeError: Batch norm got a batched tensor as # input while the running_mean or running_var, which will be updated in # place, were not batched. xfail("native_batch_norm"), })) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float,)) @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)}) @opsToleranceOverride('TestOperators', 'test_vmapjvpvjp', ( tol1('linalg.svd', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('linalg.householder_product', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('linalg.multi_dot', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), tol1('svd', {torch.float32: tol(atol=5e-04, rtol=5e-04)}), )) def test_vmapjvpvjp(self, device, dtype, op): # Since we test `jvpvjp` seperately, # in this we just check that vmap of `jvpvjp` # is correct. if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return samples = op.sample_inputs(device, dtype, requires_grad=True) # TODO: test in-place if is_inplace(op, op.get_op()): self.skipTest("Skipped! NYI: inplace-testing not supported.") return for sample in samples: fn, primals = normalize_op_input_output(op, sample) result = fn(*primals) cotangents = tree_map(lambda x: torch.randn_like(x), result) primals_tangents = tree_map(lambda x: torch.randn_like(x), primals) cotangents_tangents = tree_map(lambda x: torch.randn_like(x), cotangents) def push_vjp(primals, cotangents): _, vjp_fn = vjp(fn, *primals) return vjp_fn(cotangents) args, spec = tree_flatten(((primals, cotangents), (primals_tangents, cotangents_tangents))) def jvp_of_vjp(*args): (primals, tangents) = tree_unflatten(args, spec) primals_out, tangents_out = jvp(push_vjp, primals, tangents) flat_primals_out, _ = tree_flatten(primals_out) flat_tangents_out, _ = tree_flatten(tangents_out) return tuple(flat_primals_out + flat_tangents_out) is_batch_norm_and_training = is_batch_norm_training(op, sample.kwargs) generator = get_fallback_and_vmap_exhaustive( jvp_of_vjp, args, {}, is_batch_norm_and_training=is_batch_norm_and_training) for loop_out, batched_out in generator: self.assertEqual(loop_out, batched_out) def _make_extremal_inputs(self, shape, device): if shape is None: return (None,) return ( torch.full(shape, -1000., device=device), torch.zeros(shape, device=device), torch.full(shape, 1000., device=device), ) def _arg_and_kwarg_options(self, args_options, kwargs_options): return itertools.product(*args_options, kwargs_options) def test_extremal_numerics_nll_loss(self, device): N, C = 3, 4 d1, d2, d3 = 5, 6, 7 shapes = ( ((N, C), (N,), (C,)), ((N, C), (N,), None), ((N, C, d1, d2, d3), (N, d1, d2, d3), (C,)), ((N, C, d1, d2, d3), (N, d1, d2, d3), None), ) kwargs_options = ({'ignore_index': 0, 'reduction': 'mean'}, {'reduction': 'sum'}, {'reduction': 'none'}, {}) for input_shape, target_shape, weight_shape in shapes: input_options = self._make_extremal_inputs(input_shape, device) for input, kwargs in self._arg_and_kwarg_options((input_options,), kwargs_options): if weight_shape is None: weight = None else: weight = torch.randn(weight_shape, device=device) target = torch.randint(0, C, target_shape, device=device) target[0] = 1 # since we're ignoring index 0, at least one element must be non-zero fn = functools.partial(torch.nn.functional.nll_loss, target=target, weight=weight, **kwargs) result = fn(input) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(fn, (cotangents, input)) def test_extremal_numerics_l1_loss(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) kwargs_options = ({'reduction': 'sum'}, {'reduction': 'none'}, {}) for shape in shapes: input_options = self._make_extremal_inputs(shape, device) target_options = self._make_extremal_inputs(shape, device) for input, target, kwargs in self._arg_and_kwarg_options((input_options, target_options), kwargs_options): result = torch.nn.functional.l1_loss(input, target) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(torch.nn.functional.l1_loss, (cotangents, input, target)) def test_extremal_numerics_mse_loss(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) kwargs_options = ({'reduction': 'sum'}, {'reduction': 'none'}, {}) for shape in shapes: input_options = self._make_extremal_inputs(shape, device) target_options = self._make_extremal_inputs(shape, device) for input, target, kwargs in self._arg_and_kwarg_options((input_options, target_options), kwargs_options): result = torch.nn.functional.mse_loss(input, target) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(torch.nn.functional.mse_loss, (cotangents, input, target)) def test_extremal_numerics_softmax(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) kwargs_options = ({'dim': 1}, {}) for shape in shapes: input_options = self._make_extremal_inputs(shape, device) for input, kwargs in self._arg_and_kwarg_options((input_options,), kwargs_options): result = torch.nn.functional.softmax(input) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(torch.nn.functional.softmax, (cotangents, input)) def test_extremal_numerics_log_softmax(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) kwargs_options = ({'dim': 1}, {}) for shape in shapes: input_options = self._make_extremal_inputs(shape, device) for input, kwargs in self._arg_and_kwarg_options((input_options,), kwargs_options): result = torch.nn.functional.log_softmax(input) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(torch.nn.functional.log_softmax, (cotangents, input)) def test_extremal_numerics_cross_entropy(self, device): N, C = 3, 4 d1, d2, d3 = 5, 6, 7 shapes = ( ((N, C), (N,), (C,)), ((N, C), (N,), None), ((N, C), (N, C), (C,)), ((N, C), (N, C), None), ((C,), (), (C,)), ((C,), (), None), ((C,), (C,), (C,)), ((C,), (C,), None), ((N, C, d1, d2, d3), (N, d1, d2, d3), (C,)), ((N, C, d1, d2, d3), (N, d1, d2, d3), None), ((N, C, d1, d2, d3), (N, C, d1, d2, d3), (C,)), ((N, C, d1, d2, d3), (N, C, d1, d2, d3), None), ) for input_shape, target_shape, weight_shape in shapes: input_options = self._make_extremal_inputs(input_shape, device) kwargs_options = [{'reduction': 'sum'}, {'reduction': 'none'}, {}] if input_shape != target_shape: kwargs_options.append({'ignore_index': 0, 'reduction': 'mean'}) for input, kwargs in self._arg_and_kwarg_options((input_options,), kwargs_options): if weight_shape is None: weight = None else: weight = torch.randn(weight_shape, device=device) if input_shape == target_shape: target = torch.rand(target_shape, device=device) elif len(target_shape) == 0: target = torch.tensor(1, device=device) # must be non-zero since ignore_index may be 0 else: target = torch.randint(0, C, target_shape, device=device) fn = functools.partial(torch.nn.functional.cross_entropy, target=target, weight=weight, **kwargs) result = fn(input) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(fn, (cotangents, input), atol_rtol=(1e-4, 1e-5)) def test_extremal_numerics_binary_cross_entropy(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) for shape in shapes: weight_options = self._make_extremal_inputs(shape, device) kwargs_options = [{'reduction': 'sum'}, {'reduction': 'none'}, {}] for weight, kwargs in self._arg_and_kwarg_options((weight_options,), kwargs_options): input = torch.rand(shape, device=device) target = torch.rand(shape, device=device) fn = functools.partial(torch.nn.functional.binary_cross_entropy, target=target, weight=weight, **kwargs) result = fn(input) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(fn, (cotangents, input), atol_rtol=(1e-4, 2e-5)) def test_extremal_numerics_layer_norm(self, device): N, C, H, W = 3, 4, 5, 6 shapes = ((N, C), (N, C, H), (N, C, H, W)) for shape in shapes: input_options = self._make_extremal_inputs(shape, device) normalized_shape = shape[1:] weight_options = self._make_extremal_inputs(normalized_shape, device) bias_options = self._make_extremal_inputs(normalized_shape, device) for input, bias, weight in self._arg_and_kwarg_options((input_options, bias_options, weight_options), ()): def fn(input, weight, bias): return torch.nn.functional.layer_norm(input, normalized_shape, weight=weight, bias=bias) result = fn(input, weight, bias) cotangents = torch.randn_like(result, device=device) self._compare_jacobians_of_vjp(fn, (cotangents, input, weight, bias)) @ops(op_db + additional_op_db, allowed_dtypes=(torch.float32, torch.double)) @skipOps('TestOperators', 'test_vmap_autograd_grad', { xfail('linalg.eig'), # all close? # The size of tensor a (4) must match the size of tensor b (10) at non-singleton dimension 0 xfail('masked_select'), xfail('nn.functional.max_unpool2d', 'grad'), # contiguous call xfail('nn.functional.max_unpool2d'), # contiguous call xfail('to_sparse'), # dispatch key issue # numerical inconsistencies, look like bugs skip('matrix_exp', dtypes=(torch.float32,), device_type='cuda'), # fails on linux, passes on windows skip('ldexp', dtypes=(torch.float32,), device_type='cpu'), # fails on all but mac skip('__rmatmul__'), # flaky needs investigation skip('matmul'), # flaky needs investigation skip('nn.functional.conv_transpose3d'), # flaky needs investigation skip('nn.functional.conv_transpose2d'), # flaky needs investigation skip('nn.functional.conv_transpose1d'), # flaky needs investigation skip('nn.functional.layer_norm', dtypes=(torch.float32,), device_type='cpu'), # fails on windows skip('linalg.lu_factor', dtypes=(torch.float32,), device_type='cuda'), # fails on all but windows skip('linalg.lu_factor_ex', dtypes=(torch.float32,), device_type='cuda'), # fails on all but windows skip('linalg.multi_dot', '', device_type='cpu'), skip('sparse.sampled_addmm', ''), skip('native_layer_norm', '', device_type='cpu'), xfail('as_strided_scatter', ''), }) @opsToleranceOverride('TestOperators', 'test_vmap_autograd_grad', ( tol1('linalg.householder_product', {torch.float32: tol(atol=5e-04, rtol=9e-03)}, device_type='cuda'), tol1('linalg.householder_product', {torch.float32: tol(atol=1e-04, rtol=1e-04)}, device_type='cpu'), )) def test_vmap_autograd_grad(self, device, dtype, op): def is_differentiable(inp): return isinstance(inp, Tensor) and (inp.grad_fn is not None or inp.requires_grad) def get_flat_differentiable(pytree): flattened = tree_flatten(pytree)[0] return tuple(i for i in flattened if is_differentiable(i)) def get_differentiable_linked(list1, list2): paired_list = zip(list1, list2) paired_list = tuple((first, second) for (first, second) in paired_list if is_differentiable(first)) return zip(*paired_list) def filter_none(out): flattened = tree_flatten(out)[0] return tuple(o for o in flattened if o is not None) if not op.supports_autograd: self.skipTest("Skipped! Autograd not supported.") return sample_inputs = op.sample_inputs(device, dtype, requires_grad=True) for sample_input in sample_inputs: fn, primals = normalize_op_input_output(op, sample_input) out = fn(*primals) cotangents = tree_map(torch.randn_like, out) def compute_grad(cotangents): out_flattened = out cotangents_flattened = cotangents if not isinstance(out_flattened, torch.Tensor): out_flattened = tree_flatten(out)[0] cotangents_flattened = tree_flatten(cotangents)[0] out_flattened, cotangents_flattened = get_differentiable_linked(out_flattened, cotangents_flattened) return filter_none( torch.autograd.grad(out_flattened, get_flat_differentiable(primals), cotangents_flattened, retain_graph=True, allow_unused=True)) is_batch_norm_and_training = is_batch_norm_training(op, sample_input.kwargs) generator = get_fallback_and_vmap_exhaustive( compute_grad, (cotangents,), {}, is_batch_norm_and_training=is_batch_norm_and_training) for loop_out, batched_out in generator: self.assertEqual(loop_out, batched_out) def test_vmapvmapjvp_linalg_solve(self): ops = [op for op in op_db if op.name == "linalg.solve"] assert len(ops) > 0 # this specializes a lot of code from the get_fallback_and_vmap_exhaustive test. If we need this more # generally, this could go for a refactor B0 = 2 B1 = 3 # we want to check the case where A will be seen as contiguous by jvp but during the vmap calls will become # non-contiguous because vmap will expand. This will happen during both levels of vmap A = torch.randn(4, 4) k = torch.randn(4, 5, B1, B0) fn, args = get_jvp_variant_primals_tangents(torch.linalg.solve, SampleInput(A, args=(k,))) in_dims_all = (None, -1, None, -1) batched_out = vmap(vmap(fn, in_dims=in_dims_all), in_dims=in_dims_all)(*args) loop_out = loop2(fn, in_dims_all, in_dims_all, 0, 0, B0, B1, *args) self.assertEqual(loop_out, batched_out) @ops(filter(lambda op: op.name in aliasing_ops, op_db + additional_op_db), allowed_dtypes=(torch.float,)) @parametrize("grad_op", ["jvp", "vjp"]) def test_view_then_inplace(self, device, dtype, op, grad_op): for sample_input in op.sample_inputs(device, dtype): def f(x): op(sample_input.input, *sample_input.args, **sample_input.kwargs).copy_(x) return x without_grad = op(sample_input.input, *sample_input.args, **sample_input.kwargs) if grad_op == "jvp": with self.assertRaisesRegex(RuntimeError, "During a grad .* attempted to call in-place operation"): jvp(f, (torch.randn_like(without_grad),), (torch.randn_like(without_grad),)) else: assert grad_op == "vjp" with self.assertRaisesRegex(RuntimeError, "During a grad .* attempted to call in-place operation"): vjp(f, torch.randn_like(without_grad)) @ops(filter(lambda op: op.name in aliasing_ops_list_return, op_db + additional_op_db), allowed_dtypes=(torch.float,)) @parametrize("grad_op", ["jvp", "vjp"]) def test_view_then_inplace_list_return(self, device, dtype, op, grad_op): for sample_input in op.sample_inputs(device, dtype): def f(x): op(sample_input.input, *sample_input.args, **sample_input.kwargs)[0].copy_(x) return x without_grad = op(sample_input.input, *sample_input.args, **sample_input.kwargs)[0] with self.assertRaisesRegex(RuntimeError, "During a grad .* attempted to call in-place operation"): if grad_op == "jvp": jvp(f, (torch.randn_like(without_grad),), (torch.randn_like(without_grad),)) else: assert grad_op == "vjp" vjp(f, torch.randn_like(without_grad)) @parametrize("grad_op", ["jvp", "vjp"]) def test_view_then_inplace_special(self, grad_op): # some things in __getitem__ use at::index, which doesn't alias, so this tests a subset of them that do alias ops = [ lambda x: x[0], lambda x: x[0, 0, 0], lambda x: x[:1], lambda x: x[:, :1], lambda x: x[:, :1, :], ] for op in ops: def f(x): op(captured).copy_(x) return x captured = torch.randn(4, 3, 3) without_grad = op(captured) if grad_op == "jvp": with self.assertRaisesRegex(RuntimeError, "During a grad .* attempted to call in-place operation"): jvp(f, (torch.randn_like(without_grad),), (torch.randn_like(without_grad),)) else: assert grad_op == "vjp" with self.assertRaisesRegex(RuntimeError, "During a grad .* attempted to call in-place operation"): vjp(f, torch.randn_like(without_grad)) only_for = ("cpu", "cuda") instantiate_device_type_tests(TestOperators, globals(), only_for=only_for) if __name__ == '__main__': run_tests()