megengine.functional.argmax

argmax(inp, axis=None, keepdims=False)[source]

Returns the indices of the maximum values along given axis. If axis is a list of dimensions, reduce over all of them.

Parameters

  • inp (Tensor) – input tensor.

  • axis (Union[int, Sequence[int], None]) – dimension to reduce. If None, all dimensions will be reduced. Default: None

  • keepdims (bool) – whether the output tensor has axis retained or not. Default: False

Return type

Tensor

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)

previous

megengine.functional.argmin

next

megengine.functional.sort