Function sort

imperative/python/megengine/functional/math.py:613–652

r"""Returns sorted tensor and the indices would sort the input tensor. Args: inp: input tensor. The result would be sorted along the given dimension. descending: sort in descending order, where the largest comes first. Default: False dim: which dimension to sort along. D

(inp: Tensor, descending: bool = False, dim: int = -1)

def sort(inp: Tensor, descending: bool = False, dim: int = -1) -> Tuple[Tensor, Tensor]:
    r"""Returns sorted tensor and the indices would sort the input tensor.

    Args:
        inp: input tensor. The result would be sorted along the given dimension.
        descending: sort in descending order, where the largest comes first. Default: False
        dim: which dimension to sort along. Default: -1 (last dimension)

    Returns:
        tuple of two tensors `(sorted_tensor, indices_of_int32)`.

    Examples:
        >>> import numpy as np
        >>> x = Tensor(np.array([1,2], dtype=np.float32))
        >>> out, indices = F.sort(x)
        >>> out.numpy()
        array([1., 2.], dtype=float32)
    """
    if descending:
        order = "descending"
    else:
        order = "ascending"

    op = builtin.Argsort(order=order)

    transposed = swapaxes(inp, dim, -1)
    transposed_shape = transposed.shape
    nrows = transposed_shape[-1]
    ncols = 1
    for i in range(len(transposed_shape) - 1):
        ncols *= transposed_shape[i]
    reshaped = transposed.reshape(ncols, nrows)
    tns, ind = apply(op, reshaped)

    def restore(tensor):
        unreshaped = tensor.reshape(transposed_shape)
        untransposed = swapaxes(unreshaped, dim, -1)
        return untransposed

    return restore(tns), restore(ind)


def topk(