Contents:
Returns the unraveled index of the largest entry (by magnitude) in an n-d array
x (jax.numpy.ndarray) – An nd-array
n-dim index of the largest entry in x
tuple