def L2Loss():
  r"""Returns a layer that computes an L2-like loss for one batch.

  The layer takes three inputs:

    - Model output from one batch, an ndarray of float-valued elements.

    - A batch of element-wise target values, which matches the shape of the
      model output.

    - A batch of weights, which matches the shape of the model output.

  The layer returns a weighted average of element-wise squared error terms
  :math:`(y_i - t_i)^2`.
  """
  def f(model_output, targets, weights):  # pylint: disable=invalid-name
    shapes.assert_same_shape(model_output, targets)
    shapes.assert_same_shape(model_output, weights)
    weighted_sse = weights * (model_output - targets)**2
    return jnp.sum(weighted_sse) / jnp.sum(weights)
  return base.Fn('L2Loss', f)
Returns a layer that computes an L2-like loss for one batch.

  The layer takes three inputs:

    - Model output from one batch, an ndarray of float-valued elements.

    - A batch of element-wise target values, which matches the shape of the
      model output.

    - A batch of weights, which matches the shape of the model output.

  The layer returns a weighted average of element-wise squared error terms
  :math:`(y_i - t_i)^2`.
  
trax/layers/metrics.py