Repository: pytorch
Issue: Pass backward flags such as retain_graph to context of custom torch.autograd.Function
Labels: module: autograd,triaged

## 🚀 Feature
Write global backward flags (at least `retain_graph`) to context of custom torch.autograd.Function

## Motivation
All meta information in custom backward function is stored inside context variable typically called `ctx`. However, information about flags, which were used in backward call is not stored there. In some cases such information (mainly `retain_graph` flag) can be very crucial, for example, when another grad call is performed inside `backward` method of custom function. This is a typical approach at least in [deep implicit networks](http://implicit-layers-tutorial.org/). For example, this dummy code
```
import torch as th

class Dummy(th.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        with th.enable_grad():
            y = x**2
        ctx.save_for_backward(y, x)
        print(f'Fwd:\n{dir(ctx)}')
        return x
    
    @staticmethod
    def backward(ctx, g):
        y, x = ctx.saved_tensors
        print(f'Bkwd:\n{dir(ctx)}')
        grad_x = th.autograd.grad(y, x, g, retain_graph=False, 
                                  create_graph=False, only_inputs=True, 
                                  allow_unused=False)
        return grad_x
    
t = th.rand(5, requires_grad=True)
loss = Dummy.apply(t).sum()
loss.backward(retain_graph=True)
loss.backward(retain_graph=False)
```
fails at the second backward call: 
```
Fwd:
['__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_do_backward', '_forward_cls', '_is_legacy', '_register_hook', '_register_hook_dict', 'apply', 'dirty_tensors', 'mark_dirty', 'mark_non_differentiable', 'mark_shared_storage', 'materialize_grads', 'metadata', 'name', 'needs_input_grad', 'next_functions', 'non_differentiable', 'register_hook', 'requires_grad', 'save_for_backward', 'saved_tensors', 'saved_variables', 'set_materialize_grads', 'to_save']
Bkwd:
['__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_do_backward', '_forward_cls', '_is_legacy', '_register_hook', '_register_hook_dict', 'apply', 'dirty_tensors', 'mark_dirty', 'mark_non_differentiable', 'mark_shared_storage', 'materialize_grads', 'metadata', 'name', 'needs_input_grad', 'next_functions', 'non_differentiable', 'register_hook', 'requires_grad', 'save_for_backward', 'saved_tensors', 'saved_variables', 'set_materialize_grads', 'to_save']
Bkwd:
['__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '__weakref__', '_do_backward', '_forward_cls', '_is_legacy', '_register_hook', '_register_hook_dict', 'apply', 'dirty_tensors', 'mark_dirty', 'mark_non_differentiable', 'mark_shared_storage', 'materialize_grads', 'metadata', 'name', 'needs_input_grad', 'next_functions', 'non_differentiable', 'register_hook', 'requires_grad', 'save_for_backward', 'saved_tensors', 'saved_variables', 'set_materialize_grads', 'to_save']
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-128-f9b69f3c538c> in <module>()
     22 loss = Dummy.apply(t).sum()
     23 loss.backward(retain_graph=True)
---> 24 loss.backward(retain_graph=False)

~/anaconda3/lib/python3.6/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    243                 create_graph=create_graph,
    244                 inputs=inputs)
--> 245         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    246 
    247     def register_hook(self, hook):

~/anaconda3/lib/python3.6/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    145     Variable._execution_engine.run_backward(
    146         tensors, grad_tensors_, retain_graph, create_graph, inputs,
--> 147         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
    148 
    149 

~/anaconda3/lib/python3.6/site-packages/torch/autograd/function.py in apply(self, *args)
     87     def apply(self, *args):
     88         # _forward_cls is defined by derived class
---> 89         return self._forward_cls.backward(self, *args)  # type: ignore
     90 
     91 

<ipython-input-128-f9b69f3c538c> in backward(ctx, g)
     16         grad_x = th.autograd.grad(y, x, g, retain_graph=False, 
     17                                   create_graph=False, only_inputs=True,
---> 18                                   allow_unused=False)
     19         return grad_x
     20 

~/anaconda3/lib/python3.6/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    223     return Variable._execution_engine.run_backward(
    224         outputs, grad_outputs_, retain_graph, create_graph,
--> 225         inputs, allow_unused, accumulate_grad=False)
    226 
    227 

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling .backward() or autograd.grad() the first time.
```
As you can also see from this example is that `dir(ctx)` does not contain any global backward flags. Of course, in most real case scenarios custom function is not called like this, however there is at least one useful case, where it does - torch.autograd.gradcheck. Currently in order to test this function using gradcheck, user has to manually replace `retain_graph=False` to `retain_graph=True`, pass test and then return `retain_graph=False` back. This is not supportive if you wnat to do automative testing.
## Pitch
Before custom function backward call, write global backward flags such as `retain_graph` inside context argument. In this case, example from above can be fixed easily:
```
import torch as th

class Dummy(th.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        with th.enable_grad():
            y = x**2
        ctx.save_for_backward(y, x)
        print(f'Fwd:\n{dir(ctx)}')
        return x
    
    @staticmethod
    def backward(ctx, g):
        y, x = ctx.saved_tensors
        print(f'Bkwd:\n{dir(ctx)}')
        grad_x = th.autograd.grad(y, x, g, retain_graph=ctx.retain_graph, 
                                  create_graph=False, only_inputs=True, 
                                  allow_unused=False)
        return grad_x
    
t = th.rand(5, requires_grad=True)
loss = Dummy.apply(t).sum()
loss.backward(retain_graph=True)
loss.backward(retain_graph=False)
```

## Alternatives
Alternatively user can write several different custom autograd functions with different behaviors, but honestly this is not a solution. In example from above another hack is to do the following:
```
import torch as th

class Dummy(th.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        with th.enable_grad():
            y = x**2
            ctx.y_with_grad = True
        ctx.save_for_backward(y, x)
        print(f'Fwd:\n{dir(ctx)}')
        return x
    
    @staticmethod
    def backward(ctx, g):
        y, x = ctx.saved_tensors
        print(f'Bkwd:\n{dir(ctx)}')
        if not ctx.y_with_grad:
            with th.enable_grad():
                y = x**2
        grad_x = th.autograd.grad(y, x, g, retain_graph=False, 
                                  create_graph=False, only_inputs=True, 
                                  allow_unused=False)
        ctx.y_with_grad = False
        return grad_x
    
t = th.rand(5, requires_grad=True)
loss = Dummy.apply(t).sum()
loss.backward(retain_graph=True)
loss.backward(retain_graph=False)
```
however, it requires additional forward call inside custom backward function, hence increases runtime.  

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7