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)” and jax.nn.gelu for details.

Return type

Tuple[InitFn, ApplyFn, LayerKernelFn]

Returns

(init_fn, apply_fn, kernel_fn).