flax.linen.GRUCell#

class flax.linen.GRUCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

GRU cell.

The mathematical definition of the cell is as follows

\[\begin{split}\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h) \\ n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\ h' = (1 - z) * n + z * h \\ \end{array}\end{split}\]

where x is the input and h, is the output of the previous time step.

features#

number of output features.

gate_fn#

activation function used for gates (default: sigmoid).

activation_fn#

activation function used for output and memory update (default: tanh).

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

bias_init#

initializer for the bias parameters (default: initializers.zeros_init())

dtype#

the dtype of the computation (default: None).

param_dtype#

the dtype passed to parameter initializers (default: float32).

Parameters: