mars.tensor.argmax(a, axis=None, out=None, combine_size=None)[source]#

Returns the indices of the maximum values along an axis.

  • a (array_like) – Input tensor.

  • axis (int, optional) – By default, the index is into the flattened tensor, otherwise along the specified axis.

  • out (Tensor, optional) – If provided, the result will be inserted into this tensor. It should be of the appropriate shape and dtype.

  • combine_size (int, optional) – The number of chunks to combine.


index_array – Tensor of indices into the tensor. It has the same shape as a.shape with the dimension along axis removed.

Return type

Tensor of ints

See also

Tensor.argmax, argmin


The maximum value along a given axis.


Convert a flat index into an index tuple.


In case of multiple occurrences of the maximum values, the indices corresponding to the first occurrence are returned.


>>> import mars.tensor as mt
>>> a = mt.arange(6).reshape(2,3)
>>> a.execute()
array([[0, 1, 2],
       [3, 4, 5]])
>>> mt.argmax(a).execute()
>>> mt.argmax(a, axis=0).execute()
array([1, 1, 1])
>>> mt.argmax(a, axis=1).execute()
array([2, 2])

Indexes of the maximal elements of a N-dimensional tensor:

>>> ind = mt.unravel_index(mt.argmax(a, axis=None), a.shape)
>>> ind.execute()
(1, 2)
>>> a[ind].execute()  # TODO(jisheng): accomplish when fancy index on tensor is supported
>>> b = mt.arange(6)
>>> b[1] = 5
>>> b.execute()
array([0, 5, 2, 3, 4, 5])
>>> mt.argmax(b).execute()  # Only the first occurrence is returned.