def FastGelu():
  r"""Returns a layer that computes a fast approximation to `Gelu`.

  .. math::
      f(x) = \frac{x}{2} \cdot (1 + \tanh(ax + abx^3))

  where :math:`a = 0.7978845608` and :math:`b = 0.044715`.
  """
  def f(x):  # pylint: disable=invalid-name
    return 0.5 * x * (1 + jnp.tanh(x * 0.7978845608 * (1 + 0.044715 * x * x)))
  return Fn('FastGelu', f)
Returns a layer that computes a fast approximation to `Gelu`.

  .. math::
      f(x) = \frac{x}{2} \cdot (1 + \tanh(ax + abx^3))

  where :math:`a = 0.7978845608` and :math:`b = 0.044715`.
  
trax/layers/activation_fns.py