cruvadom:

I find the sort order of identical items to be non-deterministic (probably on GPU only).


Maybe…
# 1. get nonzero mask
s_mask = x != 0

# 2. get head mask
d_mask = torch.sum(s_mask, dim=1, keepdims=True) > torch.arange(x.shape[1])

# copy
y = torch.zeros_like(x) 
y[d_mask] = x[s_mask]