Repository: pytorch
Issue: Marking variable not required in backward calculation if it is not needed
Labels: feature,module: autograd,triaged,actionable

## 🚀 Feature

I wonder if there is a feature in pytorch's `autograd.Function` where it can detect which variables' grads are required and not required.

For example:

    import torch

    class Mul(torch.autograd.Function):
        @staticmethod
        def forward(ctx, inp1, inp2):
            ctx.save_for_backward(inp1, inp2)
            return inp1 * inp2
        @staticmethod
        def backward(ctx, grad_out):
            inp1, inp2 = ctx.saved_tensors
            grad_inp1 = None
            grad_inp2 = None
            if inp1.requires_grad:  # and inp1.grad_is_needed:
                print("Calculating grad_inp1")
                grad_inp1 = grad_out * inp2
            if inp2.requires_grad:  # and inp2.grad_is_needed:
                print("Calculating grad_inp2")
                grad_inp2 = grad_out * inp1
            return grad_inp1, grad_inp2

    dtype = torch.float64
    inp1 = torch.tensor(1.0, dtype=dtype).requires_grad_()
    inp2 = torch.tensor(2.0, dtype=dtype).requires_grad_()

    out = Mul.apply(inp1, inp2)
    # grad_inp2 is not required
    gradinp1 = torch.autograd.grad(out, inp1, create_graph=True)
    # prints:
    # Calculating grad_inp1
    # Calculating grad_inp2

Gradient of `inp2` is not required in the grad calculation, but there is no way to know if `inp2` grad is required or not.
It would be good if there is a function that can be called to see if `inp2` grad is required, e.g. `inp2.grad_is_needed`.

## Motivation

This problem comes up when I am doing multiple gradients (higher order grad) where the gradient calculation spends a lot of time in unnecessary part of the graph. Similar to the reason in #39784.

## Pitch

This feature would enable optimizing the gradient calculation for large graph, especially for computation that requires higher order gradients.

## Alternatives

For this simple case, I can just mark `inp2` as not requiring grad, but for this example below, there is no way other than just accepting it (i.e. the fact that it calculates unnecessary gradients).

    import torch

    class Pow(torch.autograd.Function):
        @staticmethod
        def forward(ctx, base, pow):
            out = base ** pow
            ctx.save_for_backward(base, pow, out)
            return out
        @staticmethod
        def backward(ctx, grad_out):
            base, pow, out = ctx.saved_tensors
            grad_base = None
            grad_pow = None
            if base.requires_grad:  # and base.grad_is_needed:
                print("Calculating grad_base")
                grad_base = grad_out * pow * out / base
            if pow.requires_grad:  # and pow.grad_is_needed:
                print("Calculating grad_pow")
                grad_pow = grad_out * out * torch.log(base)
            return grad_base, grad_pow

    dtype = torch.float64
    base = torch.tensor(1.0, dtype=dtype).requires_grad_()
    pow = torch.tensor(2.0, dtype=dtype).requires_grad_()

    out = Pow.apply(base, pow)
    gp = torch.autograd.grad(out, pow, create_graph=True)
    # calculating grad_pow is not required below
    gbgp = torch.autograd.grad(gp, base, create_graph=True)

## Additional context

This is similar to issue #39784 where it has been solved since February, and I am using nightly version `1.9.0.dev20210409` to run the example above.

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