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

  The layer takes three 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\}`, where :math:`N` is the activation vector
      depth/dimensionality.

    - A batch of weights, which matches or can be broadcast to match the shape
      of the target ndarray. This arg can give uneven weighting to different
      items in the batch (depending, for instance, on the item's target
      category).

  The predicted category from each vector is the index of the highest-valued
  vector component. The layer returns a weighted average accuracy of these
  predictions.
  """
  def f(model_output, targets, weights):  # pylint: disable=invalid-name
    predictions = jnp.argmax(model_output, axis=-1)
    shapes.assert_same_shape(predictions, targets)
    ones_and_zeros = jnp.equal(predictions, targets)
    return jnp.sum(ones_and_zeros * weights) / _n_weights_per_core(weights)

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

  The layer takes three 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\}`, where :math:`N` is the activation vector
      depth/dimensionality.

    - A batch of weights, which matches or can be broadcast to match the shape
      of the target ndarray. This arg can give uneven weighting to different
      items in the batch (depending, for instance, on the item's target
      category).

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