code: def _Accuracy():
      """Returns a layer that scores predicted versus target category."""
  
      def f(
          predicted_category, target_category
      ):  # pylint: disable=invalid-name
          # TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
          # shapes.assert_same_shape(predicted_category, target_category)
          return jnp.equal(predicted_category, target_category).astype(
              jnp.float32
          )
  
      return base.Fn("_Accuracy", 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 21 for calling a global function with example usage: _Accuracy() and returns: base, jnp