Behaviour of autograd with broadcasted tensors

username: Manecarr
cooked: <p>Hello, by trying applying <code>autograd</code> for performing reverse-mode autodiff on multi-dimensional input arrays I found this kind of behaviour:</p>
<pre><code class="lang-auto">a = torch.tensor([[1., 2., 3.], [4., 5., 6.]], requires_grad=True)
b = torch.tensor([10., 20., 30.], requires_grad=True)
q = 3*a**3 - b**2
q.backward(torch.ones_like(a))
print(torch.all(a.grad == 9*a**2).item(), torch.all(b.grad == -2*b).item())
&gt;&gt;&gt; True False
</code></pre>
<p>Indeed the value of <code>b.grad</code> is <code>-4b</code>, as if the gradients were accumulated twice. However this is not the case as <code> b.register_post_accumulate_grad_hook(lambda p: print(f"b accumulating grad of {p}: {p.grad=}"))</code> shows that the gradients were only accumulated once.</p>
<p>On the other hand, the results are different if <code>b</code> is already broadcasted to the right shape:</p>
<pre><code class="lang-auto">a = torch.tensor([[1., 2., 3.], [4., 5., 6.]], requires_grad=True)
b = torch.tensor([[10., 20., 30.], [10., 20., 30.]], requires_grad=True)
q = 3*a**3 - b**2
q.backward(torch.ones_like(a))
print(torch.all(a.grad == 9*a**2).item(), torch.all(b.grad == -2*b).item())
&gt;&gt;&gt; True True
</code></pre>
<p>So what exactly happens here? I guess the different behaviour is due to the way in which the Jacobian-vector product is implemented in <code>autograd</code>. Could you provide some details about the implementation when using arrays with more than 1 dimension (like matrices and multi-dimensional tensors) which I could not find on the online documentation?</p>
<p>Thank you very much.</p>
<p>P.S. I run the code with PyTorch build version 2.5.1+cu124</p>

username: soulitzer
cooked: <p>You tend to want to use torch.allclose instead of == as there’s no guarantee of providing bit-for-bit equivalent results.</p>
<p>The accumulate grad hooks is called when the .grad is set, but if a tensor is used multiple times, accumulation happens before.</p>

username: Manecarr
cooked: <p>Hello, thank you for the reply.</p>
<p>However the issue here is not floating-point precision, but I think I understood the mechanism: broadcasting is adding another operation to the graph and so the Jacobian of the whole transformation has to multiply the Jacobian of the broadcasting transformation (described by <code>ExpandBackward</code>). The transpose of this Jacobian matrix multiplied by the vector of ones then gives exactly what <code>b.grad</code> shows.</p>
<p>Of course if one starts with <code>b</code> in already the shape of <code>a</code>, then broadcasting does not happen and then there is one less Jacobian in the chain and of course the gradients will be different.</p>
<p>If my intuition is correct, the question could be marked as resolved thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20">.</p>

username: soulitzer
cooked: <p>Sorry you are totally right, because you implicitly broadcasted during forward. Autograd implicitly reduces during backward and you end up with .grad values that are double than if you had pre-broadcasted.</p>