def LogSumExp(axis=-1):
  """Returns a layer that computes log(sum(exp(x))) along one tensor axis.

  Args:
    axis: Axis along which values are grouped for computing log-sum-exp.
  """
  return Fn('LogSumExp',
            lambda x: fastmath.logsumexp(x, axis=axis, keepdims=True))
Returns a layer that computes log(sum(exp(x))) along one tensor axis.

  Args:
    axis: Axis along which values are grouped for computing log-sum-exp.
  
trax/layers/core.py