jax.numpy.argmax

jax.numpy.argmax(a, axis=None, out=None, keepdims=None)

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

LAX-backend 实现的 numpy.argmax()。

原始文档字符串如下。

参数:

• a (array_like) – 输入数组。

• axis (int, optional) – 默认情况下，索引是针对展平的数组，否则沿着指定的轴进行。

• keepdims (bool, optional) – 如果设置为 True，则被缩减的轴将作为大小为 1 的维度保留在结果中。通过此选项，结果将正确地与数组进行广播。

• out (None)

返回:

index_array – 数组的索引数组。它与 a.shape 具有相同的形状，但删除了沿 axis 的维度。如果 keepdims 设置为 True，则 axis 的大小将为 1，结果数组将具有与 a.shape 相同的形状。

返回类型:

ndarray of ints