def ArgMax(axis=-1):
  """Returns a layer that calculates argmax along the given axis."""
  def f(model_output):  # pylint: disable=invalid-name
    return jnp.argmax(model_output, axis=axis)
  return Fn('ArgMax', f)
Returns a layer that calculates argmax along the given axis.
trax/layers/core.py