GRUEncoder

gru_module.GRUEncoder · inherits nnx.Module

Stacked GRU encoder, scan over time, dropout BETWEEN layers (not after the last).

Input kernel and bias init from Uniform(-1/sqrt(H), 1/sqrt(H)). The recurrent kernel uses the initializer selected by recurrent_init ("uniform" or "orthogonal"). Flax nnx.GRUCell fuses input and hidden biases into a single parameter on the input projection — pinned by test_gru_encoder_has_single_fused_bias_per_cell.

__init__(self, in_features: int, hidden_size: int, n_layers: int, dropout: float, rngs: nnx.Rngs, recurrent_init: str = 'uniform')

Parameter	Type	Default	Description
in_features	int	-	(undocumented)
hidden_size	int	-	(undocumented)
n_layers	int	-	(undocumented)
dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)
recurrent_init	str	'uniform'	(undocumented)

__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray

x: [B, T, F] -> out: [B, T, H]. Float32 throughout.

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: jnp.ndarray

MLPDecoder

gru_module.MLPDecoder · inherits nnx.Module

Linear -> ReLU -> Linear. Standard 2-layer MLP head.

__init__(self, in_features: int, hidden_size: int, out_features: int, rngs: nnx.Rngs)

Parameter	Type	Default	Description
in_features	int	-	(undocumented)
hidden_size	int	-	(undocumented)
out_features	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: jnp.ndarray

GRUNet

gru_module.GRUNet · inherits nnx.Module

Encoder -> last h hidden states (with upsample path) -> MLPDecoder.

__init__(self, *, in_features: int, encoder_hidden: int, encoder_layers: int, decoder_hidden: int, decoder_layers: int, dropout: float, h: int, input_size: int, rngs: nnx.Rngs, recurrent_init: str = 'uniform')

Parameter	Type	Default	Description
in_features	int	-	(undocumented)
encoder_hidden	int	-	(undocumented)
encoder_layers	int	-	(undocumented)
decoder_hidden	int	-	(undocumented)
decoder_layers	int	-	(undocumented)
dropout	float	-	(undocumented)
h	int	-	(undocumented)
input_size	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)
recurrent_init	str	'uniform'	(undocumented)

__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: jnp.ndarray

RECURRENT_INITS

gru_module.RECURRENT_INITS

('uniform', 'orthogonal')

uniform_one_over_sqrt_h_init

gru_module.uniform_one_over_sqrt_h_init(hidden_size: int)

Initializer that draws from Uniform(-1/sqrt(H), 1/sqrt(H)).

Parameters:

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)

Returns: _UniformInit