关于torch.sort排序后的结果可导性问题

先上结论，可导的，会将梯度回传到排序前的

import torch

import torch
import torch.nn as nn


class g(nn.Module):
    def __init__(self):
        super(g, self).__init__()
        #self.k = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=1, padding=0, bias=False)

    def forward(self, z):
        a,b =torch.sort(z, dim=-1)
        return a
        #return self.k(a)


c = 2
h = 5
w = 5
z = torch.arange(0, c * h * w).float().view(1, c, h, w)
z.requires_grad = True
k=g()
r=k(z)
print(r)
r=r.sum()
loss=(r-1)*(r-1)
loss.backward()
print("zgrad", z.grad)

tensor([[[[ 0.,  1.,  2.,  3.,  4.],
          [ 5.,  6.,  7.,  8.,  9.],
          [10., 11., 12., 13., 14.],
          [15., 16., 17., 18., 19.],
          [20., 21., 22., 23., 24.]],

         [[25., 26., 27., 28., 29.],
          [30., 31., 32., 33., 34.],
          [35., 36., 37., 38., 39.],
          [40., 41., 42., 43., 44.],
          [45., 46., 47., 48., 49.]]]], grad_fn=<SortBackward>)
zgrad tensor([[[[2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.]],

         [[2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.],
          [2448., 2448., 2448., 2448., 2448.]]]])