code: 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)
  
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 6 for calling a global function with example usage: WeightedCategoryAccuracy() and returns: base, unknown