Repository: pytorch
Issue: OpInfos disabled for batched forward grad computation
Labels: module: autograd,module: tests,triaged,module: vmap,module: forward ad

NOTE: We are not planning to do any more work on vmap in core, as it will be soon replaced by a better vmap from functorch

----

Remove the `check_batched_forward_grad=False` setting on the corresponding OpInfo to replicate:


<details>
<summary>
test_forward_mode_AD_{imag,real}_cpu_complex128 (conjugate view issue, need to investigate)
</summary>

```
======================================================================
ERROR: test_forward_mode_AD_real_cpu_complex128 (__main__.TestGradientsCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 376, in instantiated_test
    raise rte
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 371, in instantiated_test
    result = test(self, **param_kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 737, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_ops.py", line 672, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op())
  File "test/test_ops.py", line 656, in _forward_grad_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "test/test_ops.py", line 587, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "test/test_ops.py", line 559, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File "/local/pytorch5/torch/testing/_internal/common_utils.py", line 2689, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1332, in gradcheck
    return _gradcheck_helper(**args)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1351, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 786, in _test_batched_grad_forward_ad
    expected = [jvp(t) for t in tangents]
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 786, in <listcomp>
    expected = [jvp(t) for t in tangents]
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 771, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "test/test_ops.py", line 547, in fn
    output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 504, in <lambda>
    gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs),  # wrapper function for gradcheck
RuntimeError: view_as_real doesn't work on unresolved conjugated tensors.  To resolve the conjugate tensor so you can view it as real, use self.resolve_conj(); however, be warned that the resulting tensor will NOT alias the original.
```

</details>

<details>
<summary>
test_forward_mode_AD_unfold_cpu_{float64,complex128} (could be an issue with scalar inputs)
</summary>

```
======================================================================
ERROR: test_forward_mode_AD_unfold_cpu_float64 (__main__.TestGradientsCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
  File /local/pytorch5/torch/testing/_internal/common_device_type.py", line 371, in instantiated_test
    result = test(self, **param_kwargs)
  File /local/pytorch5/torch/testing/_internal/common_device_type.py", line 737, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_ops.py", line 672, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op())
  File "test/test_ops.py", line 656, in _forward_grad_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "test/test_ops.py", line 587, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "test/test_ops.py", line 559, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File /local/pytorch5/torch/testing/_internal/common_utils.py", line 2689, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File /local/pytorch5/torch/autograd/gradcheck.py", line 1332, in gradcheck
    return _gradcheck_helper(**args)
  File /local/pytorch5/torch/autograd/gradcheck.py", line 1351, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File /local/pytorch5/torch/autograd/gradcheck.py", line 790, in _test_batched_grad_forward_ad
    result = _vmap(jvp)(torch.stack(tangents))
  File /local/pytorch5/torch/_vmap_internals.py", line 271, in wrapped
    batched_outputs = func(*batched_inputs)
  File /local/pytorch5/torch/autograd/gradcheck.py", line 771, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "test/test_ops.py", line 547, in fn
    output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs)
  File /local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 504, in <lambda>
    gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs),  # wrapper function for gradcheck
  File /local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 9427, in <lambda>
    op=lambda x, *args: x.unfold(*args),
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

```

</details>

- ~test_forward_mode_AD_linalg_eigh_cpu_float64` (stack not implemented for forward AD)~
<details>
<summary>
test_forward_mode_AD_nn_functional_linear_cpu_{float64,complex128} (could be a precision issue?)
</summary>

```
Got:
tensor([[[-2.2989+0.5745j,  3.0938-1.1124j,  2.1925-3.4686j, -1.6909-0.0380j]],

        [[-1.0700-1.3572j,  0.7086+5.1560j,  3.3267+3.1176j,  0.1303+0.4241j]]],
       grad_fn=<CopySlices>)

Expected:
tensor([[-2.2989+0.5745j,  3.0938-1.1124j,  2.1925-3.4686j, -1.6909-0.0380j],
        [-1.0700-1.3572j,  0.7086+5.1560j,  3.3267+3.1176j,  0.1303+0.4241j]],
       grad_fn=<StackBackward0>)
```

</details>

<details>
<summary>
test_forward_mode_AD_einsum_cpu_complex128
</summary>

```
======================================================================
ERROR: test_forward_mode_AD_einsum_cpu_complex128 (__main__.TestGradientsCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 376, in instantiated_test
    raise rte
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 371, in instantiated_test
    result = test(self, **param_kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 737, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_ops.py", line 675, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op())
  File "test/test_ops.py", line 657, in _forward_grad_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "test/test_ops.py", line 587, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "test/test_ops.py", line 559, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File "/local/pytorch5/torch/testing/_internal/common_utils.py", line 2718, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1329, in gradcheck
    return _gradcheck_helper(**args)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1348, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 782, in _test_batched_grad_forward_ad
    expected = [jvp(t) for t in tangents]
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 782, in <listcomp>
    expected = [jvp(t) for t in tangents]
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 767, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "test/test_ops.py", line 547, in fn
    output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 504, in <lambda>
    gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs),  # wrapper function for gradcheck
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 9905, in <lambda>
    op=lambda tensors, equation: torch.einsum(equation, tensors),
  File "/local/pytorch5/torch/functional.py", line 328, in einsum
    return einsum(equation, *_operands)
  File "/local/pytorch5/torch/functional.py", line 330, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: !this_view_meta->has_fw_view()INTERNAL ASSERT FAILED at "../torch/csrc/autograd/autograd_meta.cpp":171, please report a bug to PyTorch. Expected the output of forward differentiable view operations to have the tangent have the same layout as primal
```

</details>

<details>
<summary>
test_forward_mode_AD_cat_cpu_{float64, complex128}
</summary>

```
_____________ TestFwdGradientsCPU.test_forward_mode_AD_cat_cpu_float64 _____________
Traceback (most recent call last):
  File "/pytorch/test/test_ops_fwd_gradients.py", line 61, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op(), is_inplace=False)
  File "/pytorch/test/test_ops_fwd_gradients.py", line 49, in _forward_grad_helper
    call_grad_test_helper()
  File "/pytorch/test/test_ops_fwd_gradients.py", line 46, in call_grad_test_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "/pytorch/torch/testing/_internal/common_utils.py", line 4372, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "/pytorch/torch/testing/_internal/common_utils.py", line 4341, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File "/pytorch/torch/testing/_internal/common_utils.py", line 3896, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/pytorch/torch/autograd/gradcheck.py", line 1552, in gradcheck
    return _gradcheck_helper(**args)
  File "/pytorch/torch/autograd/gradcheck.py", line 1572, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File "/pytorch/torch/autograd/gradcheck.py", line 884, in _test_batched_grad_forward_ad
    result = _vmap(jvp)(torch.stack(tangents))
  File "/pytorch/torch/_vmap_internals.py", line 223, in wrapped
    batched_outputs = func(*batched_inputs)
  File "/pytorch/torch/autograd/gradcheck.py", line 865, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "/pytorch/torch/testing/_internal/common_utils.py", line 4333, in fn
    output = op.gradcheck_wrapper(variant, *positional_args, **kwargs)
  File "/pytorch/torch/testing/_internal/opinfo/core.py", line 766, in <lambda>
    gradcheck_wrapper: Callable = lambda op, *args, **kwargs: op(*args, **kwargs)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
```

</details>

as_strided when batch dims aren't at the front of the tensor
- test_forward_mode_AD_linalg_eigvalsh_cpu_float64
- test_forward_mode_AD_lu_solve_cpu_{float64,complex128}
- test_forward_mode_AD_nn_functional_pad_circular_cpu_{float64,complex128}
- test_forward_mode_AD_linalg_multi_dot_cpu_{float64,complex128}
- test_forward_mode_AD_triangular_solve_cpu_{float64,complex128}

try to inplace a batched into non-batched tensor
- test_forward_mode_AD_matrix_exp_cpu_{float64,complex128}
- test_forward_mode_AD_lu_cpu_{float64,complex128}
- test_forward_mode_AD_linalg_householder_product_cpu_{float64,complex128}
- test_forward_mode_AD_index_add_cpu_{float64,complex128}
- test_forward_mode_AD_index_copy_cpu_{float64,complex128}
- test_forward_mode_AD_index_fill_cpu_{float64,complex128}
- test_forward_mode_AD_index_put_cpu_{float64,complex128}
- test_forward_mode_AD_masked_fill_cpu_{float64,complex128}
- test_forward_mode_AD_masked_scatter_cpu_{float64,complex128}

Try to call random operations
- test_forward_mode_AD_nn_functional_dropout_cpu_float64

<details>
<summary>
test_forward_mode_AD_fill__cpu_{float64,complex128}
</summary>

```
======================================================================
ERROR: test_forward_mode_AD_fill__cpu_complex128 (__main__.TestGradientsCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 787, in _test_batched_grad_forward_ad
    result = _vmap(jvp)(torch.stack(tangents))
  File "/local/pytorch5/torch/_vmap_internals.py", line 271, in wrapped
    batched_outputs = func(*batched_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 768, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "test/test_ops.py", line 618, in fn
    output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 559, in <lambda>
    gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs),  # wrapper function for gradcheck
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 11716, in <lambda>
    op=lambda x, scalar: torch.fill_(x.clone(), scalar),
RuntimeError: output with shape [1, 5, 5, 5] doesn't match the broadcast shape [2, 5, 5, 5]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 376, in instantiated_test
    raise rte
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 371, in instantiated_test
    result = test(self, **param_kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 737, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_ops.py", line 747, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op())
  File "test/test_ops.py", line 735, in _forward_grad_helper
    call_grad_test_helper()
  File "test/test_ops.py", line 731, in call_grad_test_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "test/test_ops.py", line 659, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "test/test_ops.py", line 630, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File "/local/pytorch5/torch/testing/_internal/common_utils.py", line 2841, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1336, in gradcheck
    return _gradcheck_helper(**args)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1357, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 790, in _test_batched_grad_forward_ad
    raise GradcheckError(
torch.autograd.gradcheck.GradcheckError: While computing batched gradients, got: output with shape [1, 5, 5, 5] doesn't match the broadcast shape [2, 5, 5, 5]


gradcheck failed while testing batched gradient computation with forward-mode AD.
This test is enabled automatically when both `check_batched_grad=True`
and `check_forward_ad=True`, but can be disabled in the following ways
dependong on how the test was invoked (via a test that calls gradcheck
directly or via an autogenerated test).

If you are adding a new operator, please file an issue and then use one of the
workarounds. The workaround depends on how your test invokes gradcheck/gradgradcheck.
If the test
- manually invokes gradcheck/gradgradcheck, then call gradcheck/gradgradcheck
  with `check_batched_forward_grad=False` as a keyword argument.
- is OpInfo-based (e.g., in test_ops.py), then modify the OpInfo for the test
  to have `check_batched_forward_grad=False`
```

</details>

<details>
<summary>
test_forward_mode_AD_squeeze_cpu_{float64,complex128}
</summary>

```
======================================================================
ERROR: test_forward_mode_AD_squeeze_cpu_float64 (__main__.TestGradientsCPU)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 371, in instantiated_test
    result = test(self, **param_kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_device_type.py", line 737, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_ops.py", line 747, in test_forward_mode_AD
    self._forward_grad_helper(device, dtype, op, op.get_op())
  File "test/test_ops.py", line 735, in _forward_grad_helper
    call_grad_test_helper()
  File "test/test_ops.py", line 731, in call_grad_test_helper
    self._grad_test_helper(device, dtype, op, variant, check_forward_ad=True, check_backward_ad=False,
  File "test/test_ops.py", line 659, in _grad_test_helper
    return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad,
  File "test/test_ops.py", line 630, in _check_helper
    self.assertTrue(gradcheck(fn, gradcheck_args,
  File "/local/pytorch5/torch/testing/_internal/common_utils.py", line 2841, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1336, in gradcheck
    return _gradcheck_helper(**args)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 1357, in _gradcheck_helper
    _test_batched_grad_forward_ad(func, tupled_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 787, in _test_batched_grad_forward_ad
    result = _vmap(jvp)(torch.stack(tangents))
  File "/local/pytorch5/torch/_vmap_internals.py", line 271, in wrapped
    batched_outputs = func(*batched_inputs)
  File "/local/pytorch5/torch/autograd/gradcheck.py", line 768, in jvp
    dual_outputs = _as_tuple(func(*inputs_with_dual))
  File "test/test_ops.py", line 618, in fn
    output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs)
  File "/local/pytorch5/torch/testing/_internal/common_methods_invocations.py", line 559, in <lambda>
    gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs),  # wrapper function for gradcheck
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
```

</details>

- In-place view ops don't quite work with batched forward grad (this is because inplace view ops don't work with batched tensors either)
  - There needs to be a way to disable batched forward grad only for the inplace case. Currently the batched forward grad is also skip for the non-inplace view!

~Trying to use forward_ad with stack that does not support it:~
- ~linalg.qr, qr~

- `quantile`, `nan_quantile`: Relies on copy_ to broadcast, but the forward AD path calls broadcast_to which does not have a batching rule in core

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @mruberry