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)
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

  
trax/layers/activation_fns.py