Torch.gather does not propagate gradients?

ozppupbg: <p>Hello,</p>
<p>I’m trying to implement the Hellinger distance as a loss function. The semantic should be like CrossEntropyLoss, where the first input is the logits and the second are the expected target indices.<br>
My implementation looks like this:</p>
<pre><code class="lang-auto">def hellinger_distance(P, target):
    Pp = torch.softmax(P, dim=1)
    # Hellinger distance between probability distributions P &amp; Q is
    # = (1/sqrt(2))*sum((sqrt(p_i) - sqrt(q_i))^2)
    # If in this formula Q is a one-hot vector, this formula becomes much simpler:
    # Let t be the target index with a one in the Q distribution and every other index is zero.
    # For an index with q_i=0, the inner sum, (sqrt(p_i) - sqrt(0))^2 becomes just p_i.
    # For the index t with q_t=1, the inner sum expands to (sqrt(p_t) - sqrt(1))^2 = p_t - 2*sqrt(p_t) + 1.
    # So the entire calculation can be reduced to a sum over the entire P, including p_t.
    # And then pick out p_t from P and calculate sum(P) - 2*sqrt(p_t) + 1 as the final inner sum.
    # What remains, is just another square root and multiplication with the constant 1/sqrt(2).

    p_t = Pp.gather(1, target.view(P.shape[0], 1), sparse_grad=False).clamp(0.0001)
    if Pp.requires_grad:
        inner_sum = Pp.sum(dim=1)
        inner_sum -= p_t.view(P.shape[0]).sqrt()*2 - 1
    else:
        # if we don't need the gradients, we can simply assume the sum of all probs after softmax is 1
        inner_sum = 2 - p_t.view(P.shape[0]).sqrt()*2 # 1 - p_t.view(inner_sum.shape).sqrt()*2 + 1

    hellinger = torch.sqrt(inner_sum.clamp(0.0001))*0.7071067811865475 # 1/sqrt(2) = 0.7071067811865475
    # collect distances as batch
    collected = torch.mean(hellinger)
    return collected

x=torch.tensor([[0.0, 100.0, 0.0, 0.0], [100.0, 0.0, 0.0, 0.0], [0.0, 0.0, 100.0, 0.0],[0.0, 0.0, 0.0, 100.0]], dtype=torch.float32, requires_grad=True)
y=torch.tensor([0, 0, 0, 0]).long()
l=hellinger_distance(x, y)
l.backward()
print(l)
print(x)
print(x.grad)
</code></pre>
<p>which gives the output:</p>
<pre><code class="lang-auto">tensor(0.7480, grad_fn=&lt;MeanBackward0&gt;)
tensor([[  0., 100.,   0.,   0.],
        [100.,   0.,   0.,   0.],
        [  0.,   0., 100.,   0.],
        [  0.,   0.,   0., 100.]], requires_grad=True)
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
</code></pre>
<p>As you can see, the gradients are all zero. I traced the gradients through the calculation and everything seems fine until the gather operation. I also tried to detach the <code>p_t</code> tensor to break the backward calculation and it did not change the other values at all.<br>
So it seems, that the gather operation does not propagate its input gradients received via <code>p_t</code> to the <code>Pp</code> tensor.</p>
<p>I’m using the version <code>2.3.0a0+40ec155e58.nv24.03</code> but I also quickly checked with <code>2.1.0+cpu</code> and got the same result.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ozppupbg">@ozppupbg</a>,</p>
<p>I don’t believe <code>torch.gather</code> has a gradient function as what would be the gradient of selecting a value from a Tensor?</p>

ozppupbg: <p>Not towards the target/selector tensor, but towards the input tensor (<code>Pp</code> in this case).</p>

AlphaBetaGamma96: <p>So, after looking at the <a href="https://pytorch.org/docs/stable/generated/torch.gather.html#torch.gather" rel="noopener nofollow ugc">docs</a> for <code>torch.gather</code> and following their example, <code>torch.gather</code> does indeed have a <code>grad_fn</code>,</p>
<pre><code class="lang-auto">t = torch.tensor([[1., 2.], [3., 4.]], requires_grad=True)
loss = torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
loss.mean().backward()
t.grad
# return tensor([[0.5000, 0.0000],
#                [0.2500, 0.2500]])
</code></pre>
<p>The reason why you don’t have a <code>grad_fn</code> is because you’ve defined your outputs as dtype <code>torch.Long</code>, which doesn’t have a gradient by definition. In the example above, if you run it with <code>dtype=torch.Long</code> it will fail.</p>

ozppupbg: <p>The problem is not the long tensor. My input tensors <code>P</code> and <code>Pp</code> are float tensors.<br>
I figured it out after some experimentation with different numbers:<br>
There appears to be a weird scaling issue between the clamping and the loss.<br>
As long as the min and max of the input are not more than some value (e.g. 10) apart and I scale the loss by 1000, it works.</p>