code: def argsort_indices(a, axis=-1):
      
      a = np.asarray(a)
      ind = list(np.ix_(*[np.arange(d) for d in a.shape]))
      ind[axis] = a.argsort(axis)
      return tuple(ind)
docs: Like argsort, but returns an index suitable for sorting the
      the original array even if that array is multidimensional