How does autograd know which gradient to pass to backward?

Andrew_Holmes: <p>I was tinkering around trying to understand how autograd works in the background. For functions that output a single tensor, its easier to interpret how the backward pass goes, but it gets more complex (IMO) when there’s multiple outputs. To get an idea of what’s happening I created a custom function that takes an input and outputs two values (the double and triple of the input respectively). According to the autograd custom <code>Function</code> <a href="https://pytorch.org/docs/stable/notes/extending.html#:~:text=It%20will%20be%20given%20as%20many%20Tensor%20arguments%20as%20there%20were%20outputs" rel="noopener nofollow ugc">rules</a>, the number of <code>grad_ouputs</code> passed to the <code>backward()</code> method should be the number of outputs (obviously). However, what I noticed is that autograd somehow knew which gradient to pass to what argument when only on the two outputs called <code>backward()</code>. To get a sense of what I mean, here’s some code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.autograd.function import Function


class DoubleTriple(Function):

    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return a * 2, a * 3

    @staticmethod
    def backward(ctx, grad1, grad2):
        print(grad1, grad2)
        return grad1 * 2 + grad2 * 3


def dt(a):
    return DoubleTriple.apply(a)


def main():

    a = torch.tensor(3.0, requires_grad=True)
    b, c = dt(a)
    b.backward()


if __name__ == "__main__":
    main()
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">tensor(1.) tensor(0.)
</code></pre>
<p>My question is how does autograd know that <code>b</code> is the first gradient given both the <code>grad_fn</code> attributes for <code>b</code> and <code>c</code> point to the same object? I tried looking more into <code>Node</code> class but there wasn’t anything I could find that indicated which gradient (implicitly passed from <code>.backward()</code>) should go to which grad argument for my custom backward function; it seems to just <em>know</em>.</p>

IpsumDominum: <p>In python if you write</p>
<pre><code class="lang-auto">args = [1,2,3]
def function(a,b,c):
      print(a,b,c)
function(*args)

Output will be : 1, 2, 3
</code></pre>
<p>Perhaps here the grads are passed in such a way?<br>
Is this what you are asking?</p>

Andrew_Holmes: <p>Some what sort of, I think this actually gives me some intuition. Maybe autograd makes a default for the <code>grad_outputs</code> and then just populates the indices to which ever gradient is provided. My question is mainly asking how autograd knows <code>b</code> is the first argument to my function and the reason why I ask, even though it seems obvious is there’s no where in their source code that I could find that shows <code>b</code> being marked as the first argument. Theoretically, it could just be storing the argument’s order in the forward pass them mapping them whenever the <code>backward()</code> method, but I myself can’t prove that. I just want to understand where this is done and how it works out of general curiosity.</p>

Andrew_Holmes: <p>Update: it’s in my face. It appears to be an attribute of the <code>Tensor</code> class that gets assigned at some point during the forward pass.</p>
<p><a href="https://github.com/pytorch/pytorch/blob/main/torch/autograd/graph.py#L173" rel="noopener nofollow ugc">torch.autograd.graph.py</a></p>
<pre data-code-wrap="python"><code class="lang-python">def get_gradient_edge(tensor):
    """Get the gradient edge for computing the gradient of the given Tensor.

    In particular, it is equivalent to call
    ``g = autograd.grad(loss, input)`` and ``g = autograd.grad(loss, get_gradient_edge(input))``.
    """
    if not tensor.requires_grad:
        raise RuntimeError(
            "It is not possible to get the gradient edge for a Tensor that does not require gradients"
        )
    grad_fn = _get_grad_fn_or_grad_acc(tensor)

    # Note that output_nr default to 0 which is the right value
    # for the AccumulateGrad node.
    return GradientEdge(grad_fn, tensor.output_nr)

</code></pre>