You can imagine sort() as doing two operations:
indices = argsort(x,dim)
y = gather(x,indices,dim) #x[indices] may work too
argsort() defines non-trainable permutation
gather() applies this permutation. It is a differentiable operation - backward pass permutes gradients back to match their original positions.
Note that you can use indices to gather from other tensors, that basically performs sortByKey() operation. So what you did, sorting z and reordering bigger tensor is valid.