Compute only the required gradients in the backward pass

[tvercaut]

Currenly if A and B are inputs, we always compute the gradient with respect to both A and B in the backward pass irrespective of whether this is requested by the caller.

As shown in the extending pytorch example
https://pytorch.org/docs/stable/notes/extending.html#example
we should make use of a pattern invoving

if ctx.needs_input_grad[i]:

to avoid unnecessary computation.