mars.tensor.argmax#

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

返回沿着某个轴的最大值的索引。

Parameters

a (array_like) – 输入张量。
axis (int, 可选) – 默认情况下，索引进入展平的张量，否则沿指定的轴。
out (张量, 可选) – 如果提供，该结果将插入到此张量中。它应该是适当的形状和数据类型。
combine_size (int, 可选) – 拼接的块数。

Returns

index_array – 索引张量。它的形状与 a.shape 相同，但去掉了沿 axis 的维度。

Return type

整数的张量

另请参阅

Tensor.argmax, argmin

amax: 在给定轴上的最大值。
unravel_index: 将平面索引转换为索引元组。

备注

如果最大值出现多个，返回对应于第一次出现的索引。

示例

>>> 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()
5
>>> mt.argmax(a, axis=0).execute()
array([1, 1, 1])
>>> mt.argmax(a, axis=1).execute()
array([2, 2])

一个N维张量的最大元素的索引：

>>> 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.
1