ArgMax

Return the index of the maximum value of an array.

Abstract Signature:

ArgMax(a: Tensor, axis: int | None, keepdims: bool = False)

PyTorch

API: torch.argmax
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: jax.numpy.argmax
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: numpy.argmax
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.argmax
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.math.argmax
Strategy: Direct Mapping
Official Docs ↗

Apple MLX

API: mlx.core.argmax
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: jax.numpy.ndarray.argmax
Strategy: Direct Mapping
Official Docs ↗