Hi,
I’m having an issue with a custom autogrid function that implements a sorting operation on an input x.
The operation sorts the rows of X based on taking the sorted indices of an input vector s.
The sorting part of the operation is all fine and working outside of the method, but I am having problems implementing it inside the method.
from torch.autograd import Function
class SortBy(Function):
@staticmethod
def forward(ctx, x, s):
_, indices = torch.sort(s)
results = x[indices, :]
ctx.save_for_backward(results)
return result
@staticmethod
def backward(ctx, grad_output, s):
results, = ctx.saved_tensors
_, indices = torch.sort(s)
_, index2 = torch.sort(indices)
return (results * grad_output)[index2, :]
I’m attempting to test my code using the following:
x = torch.autograd.Variable(torch.randn(4,5),requires_grad=True)
s = torch.randn(4)
Y = SortBy.apply(X, s)
Yhat = torch.sum(Y)
Yhat.backward(s)
But I get the below error:
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-283-b5596230715d> in <module>()
7 Yhat = torch.sum(Y)
8
----> 9 Yhat.backward(s)
2 frames
/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py in _make_grads(outputs, grads)
27 + str(grad.shape) + " and output["
28 + str(outputs.index(out)) + "] has a shape of "
---> 29 + str(out.shape) + ".")
30 new_grads.append(grad)
31 elif grad is None:
RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([4]) and output[0] has a shape of torch.Size([]).
I don’t understand where the dimension mismatch is coming in. I assume that the grad_output from the torch.sum step returns a single scalar gradient which I then broadcast to my matrix results in order to get the results, then I apply my reverse indexing to ‘reroute’ the gradient to where I want it.
