mlx.nn.glu

Contents

mlx.nn.glu#

class glu(x: array, axis: int = -1)#

Applies the gated linear unit function.

This function splits the axis dimension of the input into two halves (\(a\) and \(b\)) and applies \(a * \sigma(b)\).

\[\textrm{GLU}(x) = a * \sigma(b)\]
Parameters:

axis (int) – The dimension to split along. Default: -1