code: def _WeightedSequenceMean():
      """Returns a layer that computes a weighted sequence accuracy mean."""
  
      def f(values, weights):  # pylint: disable=invalid-name
          # This function assumes weights are 0 or 1.
          # Then compute 1: not-correct, 0: correct or masked
          not_correct = (1.0 - values) * weights
          axis_to_sum = list(range(1, len(not_correct.shape)))
          # Summing not-correct on all axes but batch. We're summing 0s and 1s,
          # so the sum is 0 if it's all 0 and >=1 in all other cases.
          not_correct_seq = jnp.sum(not_correct, axis=axis_to_sum)
          # Sequence is correct if not_correct_seq is 0, reverting here.
          correct_seq = 1.0 - jnp.minimum(1.0, not_correct_seq)
          return jnp.mean(correct_seq)  # Mean over batch.
  
      return base.Fn("_WeightedSequenceMean", 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 27 for calling a global function with example usage: _WeightedSequenceMean() and returns: base, jnp