code: def Glu():
      r"""Returns a layer that computes the Gated Linear Unit function.
  
      .. math::
          f(x) = a \cdot \text{sigmoid}(b)
      where a and b are formed by splitting input in half along axis
  
      """
  
      def _f(x, axis=-1):  # pylint: disable=invalid-name
          size = x.shape[axis]
          assert (
              size % 2 == 0
          ), f"axis {axis} of size {size} is not be divisible by 2"
          a, b = jnp.split(x, 2, axis)
          return a * fastmath.expit(b)
  
      return Fn("Glu", _f)
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 16 for calling a global function with example usage: Glu() and returns: Fn, unknown