ArgMax ====== Return the index of the maximum value of an array. **Abstract Signature:** ``ArgMax(a: Tensor, axis: int | None, keepdims: bool = False)`` .. raw:: html

PyTorch

API: torch.argmax
Strategy: Direct Mapping

JAX (Core)

API: jax.numpy.argmax
Strategy: Direct Mapping

NumPy

API: numpy.argmax
Strategy: Direct Mapping

Keras

API: keras.ops.argmax
Strategy: Direct Mapping

TensorFlow

API: tf.math.argmax
Strategy: Direct Mapping

Apple MLX

API: mlx.core.argmax
Strategy: Direct Mapping

Flax NNX

API: jax.numpy.ndarray.argmax
Strategy: Direct Mapping