neural_tangents.stax.Gelu

neural_tangents.stax.Gelu(approximate=False)[source]

Gelu function.

Parameters: approximate (bool) – only relevant for finite-width network, apply_fn. If True, computes an approximation via tanh, see https://arxiv.org/abs/1606.08415 and jax.nn.gelu for details.
Return type: Tuple[InitFn, ApplyFn, LayerKernelFn]
Returns: (init_fn, apply_fn, kernel_fn).