Are torch.sort() and torch.take() differentiable?

wgu93: <p>Hi there! I am customizing my own loss function using <code>torch.sort()</code> and <code>torch.take()</code> functions. However, with the custom loss function, a fully-connected feed-forward neural network (nothing special, with linear, ReLU and batch norm layers) does not learn well on the dataset. Therefore, I am wondering if there exist some issues with my loss function, e.g., the gradients are not back-propagated due to the use of these two functions. I checked the official document and also read some threads on Internet, but I don’t find any answers for that, unfortunately.</p>
<p>By the way, my motivation of defining my loss function is to penalize the violation of monotonicity of two series of data. For example, I got two tensors <code>x, y</code> which contain the time-series of two variables, and I knew that <code>y</code> should increase/decrease as <code>x</code> increases/decreases. Hence, I sorted<code> x</code> (in ascending order) and recorded the index (which was done by <code>torch.sort()</code>), and re-ordered <code>y</code> using that index (which was done by <code>torch.take()</code>). Lastly, I just needed to check the difference between the consecutive pair of data in sorted <code>y</code> from which monotonicity can be checked.</p>
<p>I really appreciate any feedback and thanks in advance for your time!</p>

AlphaBetaGamma96: <p>In terms of <code>torch.sort</code> or <code>torch.take</code> they do have <code>grad_fn</code>s, except for the indices returned via <code>torch.sort</code></p>
<pre><code class="lang-auto">&gt;&gt;&gt; x = torch.randn(10,2, requires_grad=True)
&gt;&gt;&gt; torch.take(x, torch.LongTensor([1]))
tensor([-0.2461], grad_fn=&lt;TakeBackward0&gt;)
&gt;&gt;&gt; torch.sort(x)[0]
tensor([[-1.2673, -0.2461],
        [-0.6003,  0.2518],
        [-1.1847,  0.2308],
        [ 0.1341,  2.5518],
        [-3.3000,  0.4909],
        [-1.2083,  1.7219],
        [-0.1799,  1.0432],
        [-1.3900,  0.2327],
        [-0.3572,  0.0238],
        [-0.7487,  0.4416]], grad_fn=&lt;SortBackward0&gt;)
&gt;&gt;&gt; torch.sort(x)[1]
tensor([[0, 1],
        [1, 0],
        [0, 1],
        [1, 0],
        [1, 0],
        [1, 0],
        [0, 1],
        [0, 1],
        [0, 1],
        [0, 1]])
</code></pre>
<p>Also, one thing to check within your loss function is to check if any intermediate gradients are 0 because gradients are calculated via the chain rule, and so if one intermediate gradient is 0 then all of them will be 0.</p>
<p>Also, check to make sure you’re not calling <code>.detach()</code> or <code>torch.no_grad()</code> within any part of your code, and make sure your data has <code>requires_grad=True</code> otherwise your gradient will be 0 by definition.</p>

KFrank: <p>Hi wgu!</p>
<aside class="quote no-group" data-username="wgu93" data-post="1" data-topic="134985" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/3bc359/48.png" class="avatar"> wgu93:</div>
<blockquote>
<p>my loss function is to penalize the violation of monotonicity of two series of data.</p>
</blockquote>
</aside>
<p>If I understand your use case correctly, when you have<br>
<code>x[i + 1] &gt; x[i]</code>, you would also like to have <code>y[i + 1] &gt; y[i]</code>,<br>
and you would like to impose a penalty if <code>y[i + 1] &lt; y[i]</code>.  (And,<br>
analogously, switching <code>&lt;</code> for <code>&gt;</code>.)</p>
<p>But you don’t care <em>how much</em> this joint monotonicity is violated.  If<br>
<code>y[i + 1]</code> is smaller than <code>y[i]</code> by just a little bit, you want to penalize<br>
it by essentially the same amount as when <code>y[i + 1]</code> is smaller by a<br>
lot.</p>
<p>Also, I assume that you want the gradient of your “non-monotonicity<br>
loss” with respect to the values of <code>x</code> and <code>y</code>.  For example, perhaps<br>
<code>x</code> and <code>y</code> are generated by a neural network, and you want to train<br>
that network to only generate such data series that are monotone<br>
with respect to one another.  Thus, you would want to backpropagate<br>
the gradient of your loss through <code>x</code> and <code>y</code> to the network parameters<br>
that generated them.</p>
<p>You don’t need (nor I think want) to use sorting for this.</p>
<p><code>torch.tanh()</code> is a differentiable (“soft”) approximation to the<br>
<code>sign()</code> function.  So you could use (in vector form):</p>
<pre data-code-wrap="python"><code class="lang-python">loss = (1.0 - (torch.tanh (alpha * (x[1:] - x[:-1])) * torch.tanh (alpha * (y[1:] - y[:1])))).mean()
</code></pre>
<p>where <code>alpha</code> can be used to sharpen (or smooth out) the abruptness<br>
with which <code>tanh()</code> approximates <code>sign()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

wgu93: <p>Thanks AlphaBetaGamma96! That was a quick check on the <code>grad_fn</code>!</p>
<p>One more question: if I use the indices returned from <code>torch.sort()</code> and feed it to <code>torch.take()</code>, will the gradients be backpropagated? (I think yes, and I posted my test code as follows.)</p>
<p>Test code:</p>
<pre><code class="lang-auto">import torch
from torch import nn
x = torch.tensor([1, 4, 7, 3, 9, 0], dtype=torch.float, requires_grad=True)
x_sorted = torch.sort(x)[0]
x_id = torch.sort(x)[1]
print(x_sorted, x_id)
y = torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.float, requires_grad=True)
y_sorted = torch.take(y, x_id) # Here is my question: Will the gradients be backpropagated?
print(y, y_sorted)

loss_mse = nn.MSELoss()
loss_xy = loss_mse(x, y)
loss_ys = y_sorted.mean()
loss_xs =  x_sorted.mean()

loss_ys.backward()
loss_xs.backward()
loss_xy.backward()
print(f"Gradients: {loss_ys.item()}, {loss_xs.item()}, {loss_xy.item()}")

loss_xid = x_id.float().mean()
loss_xid.backward()
</code></pre>
<p>which returns</p>
<pre><code class="lang-auto">tensor([0., 1., 3., 4., 7., 9.], grad_fn=&lt;SortBackward&gt;) tensor([5, 0, 3, 1, 2, 4])
tensor([0., 1., 2., 3., 4., 5.], requires_grad=True) tensor([5., 0., 3., 1., 2., 4.], grad_fn=&lt;TakeBackward&gt;)
Gradients: 2.5, 4.0, 14.166666984558105

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-30-1004213fe0c3&gt; in &lt;module&gt;
     20 
     21 loss_xid = x_id.float().mean()
---&gt; 22 loss_xid.backward()

~/miniconda3/envs/base-tornado5/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    243                 create_graph=create_graph,
    244                 inputs=inputs)
--&gt; 245         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    246 
    247     def register_hook(self, hook):

~/miniconda3/envs/base-tornado5/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    143         retain_graph = create_graph
    144 
--&gt; 145     Variable._execution_engine.run_backward(
    146         tensors, grad_tensors_, retain_graph, create_graph, inputs,
    147         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>It can be seen that gradients are calculated for <code>loss_ys</code> (which does not have any actual meaning but only takes average on <code>y_sorted</code> for testing purposes). And for <code>loss_xid</code>, error popped up due to the lack of <code>grad_fn</code> for <code>torch.sort(x)[1]</code> just like AlphaBetaGamma96 pointed out. So my answer to my question is yes: <code>torch.take(y, x_id)</code> supports gradient backpropagation with <code>x_id = torch.sort(x)[1]</code> and <code>y</code> being a  tensor that has <code>requires_grad=True</code>.</p>
<p>Let me know if I am making mistakes here. Thanks!</p>

wgu93: <p>Thanks a million, KFrank! That’s absolutely what I want for calculating the loss. I’m using it to train the network and it seems working better than my previous implementation of loss function. Very beautiful solution!</p>