def sort_1d(input):
# TODO: this isn't an ideal implementation.
val = nd.sort(input, axis=None, is_ascend=True)
idx = nd.argsort(input, is_ascend=True)
idx = nd.cast(idx, dtype='int64')
return val, idx
def argtopk(input, k, dim, descending=True):
idx = nd.argsort(input, dim, is_ascend=not descending)
return nd.slice_axis(input, dim, 0, k)
def argsort(input, dim, descending):
idx = nd.argsort(input, dim, is_ascend=not descending)
idx = nd.cast(idx, dtype='int64')
return idx
def top_k(array, k):
s = nd.argsort(array)[::-1][:k]
return array[s], s
