def _WeightedMean():
  """Returns a layer that computes a weighted mean of the given values."""
  def f(values, weights):  # pylint: disable=invalid-name
    return jnp.sum(values * weights) / _n_weights_per_core(weights)
  return base.Fn('_WeightedMean', f)
Returns a layer that computes a weighted mean of the given values.
trax/layers/metrics.py