jax.nn.glu

Contents

jax.nn.glu#

jax.nn.glu(x, axis=-1)[source]#

Gated linear unit activation function.

Computes the function:

\[\mathrm{glu}(x) = x\left[\ldots, 0:\frac{n}{2}, \ldots\right] \cdot \mathrm{sigmoid} \left( x\left[\ldots, \frac{n}{2}:n, \ldots\right] \right)\]

where the array is split into two along axis. The size of the axis dimension must be divisible by two.

Parameters:
  • x (jax.typing.ArrayLike) – input array

  • axis (int) – the axis along which the split should be computed (default: -1)

Returns:

An array.

Return type:

Array

See also

sigmoid()