def SmoothL1Loss():
  r"""Returns a layer that computes a weighted, smoothed L1 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 computes a "smooth" L1 loss (a.k.a. Huber loss), for model output
  float :math:`y_i` and target float :math:`t_i`:

  .. math::
      \text{output} = \left\{ \begin{array}{cl}
          \frac 1 2 (y_i - t_i)^2, & \text{if}\ |y_i - t_i| < 1, \\
          |y_i - t_i| - \frac 1 2, & \text{otherwise}.
      \end{array} \right.

  The layer returns a weighted average of these element-wise values.
  """
  def f(model_output, targets, weights):  # pylint: disable=invalid-name
    shapes.assert_same_shape(model_output, targets)
    shapes.assert_same_shape(model_output, weights)
    l1_dist = jnp.abs(model_output - targets)
    smooth_dist = jnp.where(l1_dist < 1, 0.5 * l1_dist**2, l1_dist - 0.5)
    weighted_smooth_dist = weights * smooth_dist
    return jnp.sum(weighted_smooth_dist) / jnp.sum(weights)
  return base.Fn('SmoothL1Loss', f)
Returns a layer that computes a weighted, smoothed L1 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 computes a "smooth" L1 loss (a.k.a. Huber loss), for model output
  float :math:`y_i` and target float :math:`t_i`:

  .. math::
      \text{output} = \left\{ \begin{array}{cl}
          \frac 1 2 (y_i - t_i)^2, & \text{if}\ |y_i - t_i| < 1, \\
          |y_i - t_i| - \frac 1 2, & \text{otherwise}.
      \end{array} \right.

  The layer returns a weighted average of these element-wise values.
  
trax/layers/metrics.py