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 “Gaussian Error Linear Units (GELUs)” andjax.nn.gelu
for details.- Return type:
- Returns:
(init_fn, apply_fn, kernel_fn).