code: def _CrossEntropy():
      """Returns a layer that computes prediction-target cross entropies."""
  
      def f(
          model_output, target_category
      ):  # pylint: disable=invalid-name
          # TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
          # shapes.assert_shape_equals(target_category, model_output.shape[:-1])
          target_distribution = core.one_hot(
              target_category, model_output.shape[-1]
          )
          return -1.0 * jnp.sum(
              model_output * target_distribution, axis=-1
          )
  
      return base.Fn("_CrossEntropy", 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 22 for calling a global function with example usage: _CrossEntropy() and returns: base, unknown