jax.lax.argmax# jax.lax.argmax(operand, axis, index_dtype)[source]# Computes the index of the maximum element along axis. Parameters: operand (jax.typing.ArrayLike) axis (int) index_dtype (jax.typing.DTypeLike) Return type: Array
