Repository: pytorch
Issue: Create a boxed fallback / template recipe for Autograd that forwards, but errors on backwards
Labels: module: internals,module: autograd,triaged,module: dispatch

FB only reference: https://fb.workplace.com/groups/pytorch.dev/posts/933475730564241/

There is a longstanding problem with custom operators to PyTorch, which is that they typically have incorrect autograd functionality. The most common practice (as documented in https://pytorch.org/tutorials/advanced/dispatcher.html) is if you write a forward-only operator, you simply omit the Autograd dispatch key definition. However, the default implementation of autograd in aten/src/ATen/core/VariableFallbackKernel.cpp simply *falls through*, so if any inputs require grad, this is silently lost. This can be **very** difficult to debug (that is what the FB reference above cites.) The correct behavior in this situation is to perform the forwards, but register the grad_fn NotImplemented (in torch/csrc/autograd/functions/basic_ops.h) which will report an error saying gradient is not implemented on backwards.

The first step to fixing this situation is making it possible to conveniently do this, in a single line. It is possible to do this by using the custom C++ autograd function interface, but it is very wordy: https://github.com/pytorch/vision/blob/b8a8fbc9933de1df3b269ac723e0dc66eefa4d80/torchvision/csrc/ops/autograd/roi_pool_kernel.cpp#L65-L102 There are two ways we can condense this:

1. Create a template function that generates all of the necessary boilerplate. As PyTorch core, we can do better than the sample code in torchvision; we should generate the equivalent of not implemented VariableType definitions (NB: I have excised the NDEBUG statements, but in a fully automated luxury space communism solution those would be implemented as well):

```
at::Tensor _add_relu_Tensor(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  auto& self_ = unpack(self, "self", 0);
  auto& other_ = unpack(other, "other", 1);
  auto _any_requires_grad = compute_requires_grad( self, other );
  (void)_any_requires_grad;
  std::shared_ptr<NotImplemented> grad_fn;
  if (_any_requires_grad) {
    grad_fn = std::shared_ptr<NotImplemented>(new NotImplemented("_add_relu"), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self, other ));
  }
  auto _tmp = ([&]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::_add_relu(ks & c10::after_autograd_keyset, self_, other_, alpha);
  })();
  auto result = std::move(_tmp);
  if (grad_fn) {
      set_history(flatten_tensor_args( result ), grad_fn);
  }   
  throw_error_for_complex_autograd(result, "_add_relu");
  TORCH_CHECK_NOT_IMPLEMENTED(!(isFwGradDefined(self) || isFwGradDefined(other)), "Trying to use forward AD with _add_relu that does not support it.");
  return result;
} 
```

This solution would be sufficient to let individual operator writers explicitly turn on correct autograd handling for their operators; but they would be required to do so explicitly. (It might be possible to modify def/impl to also automatically apply this registration when someone registers a CPU/CUDA implementation directly, but we would need to implement https://github.com/pytorch/pytorch/issues/35041 so that subsequent registrations of autograd would override this behavior.)

2. A simpler, albeit less efficient, way to solve this problem is to write the template function as a *boxed fallback*. We would thus instead write this function:

```
void autogradNotSupportedFallback(const c10::OperatorHandle& op, DispatchKeySet dispatch_keys, torch::jit::Stack* stack) {
...
}
```

This function would iterate over the stack checking for requires_grad'ness, and then set history for each of the outputs. There are many examples of existing boxed fallbacks and thus the implementation here should be relatively straightforward. The boxed fallback has an added benefit that we could then consider, at a later point in time, to register it as the universal Autograd fallback key, thus automatically marking all custom operators as not supporting autograd unless they explicitly register an Autograd implementation (this would be quite BC breaking and there would be quite a bit of fallout to clean up afterwards).

----

An orthogonal solution to this class of problems is to take @bdhirsh's work on external codegen and apply it to custom operators; in such situations, we could automatically codegen autograd kernels for custom operators. However, this is a bit further out, and I think the boxed fallback would be quickest to execute on.

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @bhosmer @smessmer @ljk53 @bdhirsh