def CategoryAccuracy():
  r"""Returns a layer that computes category prediction accuracy.

  The layer takes two inputs:

    - A batch of activation vectors. The components in a given vector should
      be mappable to a probability distribution in the following loose sense:
      within a vector, a higher component value corresponds to a higher
      probability, such that argmax within a vector (``axis=-1``) picks the
      index (category) having the highest probablity.

    - A batch of target categories; each target is an integer in
      :math:`\{0, ..., N-1\}`.

  The predicted category from each vector is the index of the highest-valued
  vector component. The layer returns the accuracy of these predictions
  averaged over the batch.
  """
  def f(model_output, targets):  # pylint: disable=invalid-name
    predictions = jnp.argmax(model_output, axis=-1)
    shapes.assert_same_shape(predictions, targets)
    n_total = predictions.size
    n_correct = jnp.sum(jnp.equal(predictions, targets))
    return n_correct / n_total

  return base.Fn('CategoryAccuracy', f)
Returns a layer that computes category prediction accuracy.

  The layer takes two inputs:

    - A batch of activation vectors. The components in a given vector should
      be mappable to a probability distribution in the following loose sense:
      within a vector, a higher component value corresponds to a higher
      probability, such that argmax within a vector (``axis=-1``) picks the
      index (category) having the highest probablity.

    - A batch of target categories; each target is an integer in
      :math:`\{0, ..., N-1\}`.

  The predicted category from each vector is the index of the highest-valued
  vector component. The layer returns the accuracy of these predictions
  averaged over the batch.
  
trax/layers/metrics.py