Function argmax

imperative/python/megengine/functional/math.py:530–574

r"""Returns the indices of the maximum values along given axis. If axis is a list of dimensions, reduce over all of them. Args: inp: input tensor. axis: dimension to reduce. If None, all dimensions will be reduced. Default: None keepdims: whether the output tenso

(
    inp: Tensor,
    axis: Optional[Union[int, Sequence[int]]] = None,
    keepdims: bool = False,
)

def argmax(
    inp: Tensor,
    axis: Optional[Union[int, Sequence[int]]] = None,
    keepdims: bool = False,
) -> Tensor:
    r"""Returns the indices of the maximum values along
    given axis. If axis is a list of dimensions,
    reduce over all of them.

    Args:
        inp: input tensor.
        axis: dimension to reduce. If None, all dimensions will be reduced. Default: None
        keepdims: whether the output tensor has axis retained or not. Default: False

    Returns:
        output tensor.

    Examples:
        >>> import numpy as np
        >>> x = Tensor(np.arange(1, 7, dtype=np.int32).reshape(2,3))
        >>> F.argmax(x)
        Tensor(5, dtype=int32, device=xpux:0)
    """
    if axis is None:
        assert not keepdims, "can not set axis=None and keepdims=True"
        inp = inp.flatten()
        axis = 0
    axis = _normalize_axis(inp.ndim, axis, reverse=True)

    if isinstance(axis, collections.abc.Iterable):

        for ai in axis:
            op = builtin.Argmax(axis=ai)
            (inp,) = apply(op, inp)

            if not keepdims:
                inp = squeeze(inp, ai)

        return inp

    op = builtin.Argmax(axis=axis)
    (result,) = apply(op, inp)
    if not keepdims:
        result = squeeze(result, axis)
    return result


# sorting functions