Returns the indices of the maximum values along an axis, ignoring NaNs.
By default, index is into the flattened array, otherwise it is along the
specified axis. Unlike NumPy, which raises an error for all-NaN slices, this
returns -1 for those slices, matching JAX.
Returns the indices of the maximum values along an axis, ignoring NaNs.
By default, index is into the flattened array, otherwise it is along the specified axis. Unlike NumPy, which raises an error for all-NaN slices, this returns -1 for those slices, matching JAX.