def LogLoss(distribution, **unused_kwargs):  # pylint: disable=invalid-name
  """Builds a log loss layer for a Distribution."""
  return tl.Serial(
      distribution.LogProb(),
      tl.Negate(),
      tl.WeightedSum()
  )
Builds a log loss layer for a Distribution.
trax/rl/distributions.py