code: 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)
  
imports: from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
label_desc: Use image node_id 1 for calling a global function with example usage: CategoryAccuracy() and returns: base, unknown