code: 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),
      )
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 4 for calling a global function with example usage: LogSumExp(axis) and returns: Fn