RNN

model.RNN · inherits flax.linen.Module

Nixtla NeuralForecast RNN ported to Flax, with GRU and LayerNorm.

Forward pass is unchanged vs v1; new config fields (cell_type, layer_norm) are threaded through to :class:RNNEncoder.

__init__(self, config)

Parameter	Type	Default	Description
config	RNNConfig	-	(undocumented)

__call__(self, insample_y, hist_exog=None, futr_exog=None, stat_exog=None, rnn_state=None, deterministic=True)

(The forward pass of the RNN model.)

Parameter	Type	Default	Description
insample_y	jnp.ndarray	-	(undocumented)
hist_exog	Optional[jnp.ndarray]	None	(undocumented)
futr_exog	Optional[jnp.ndarray]	None	(undocumented)
stat_exog	Optional[jnp.ndarray]	None	(undocumented)
rnn_state	Optional[jnp.ndarray]	None	(undocumented)
deterministic	bool	True	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray] (The output predictions [B, h, output_size] and the final RNN state [num_layers, B, H]).

RNNConfig

model.RNNConfig

Hyperparameters for :class:RNN.

Attribute	Type	Default	Description
h	int	12	(undocumented)
input_size	int	36	(undocumented)
encoder_hidden_size	int	128	(undocumented)
encoder_n_layers	int	2	(undocumented)
encoder_activation	str	"tanh"	(undocumented)
encoder_bias	bool	True	(undocumented)
encoder_dropout	float	0.0	(undocumented)
decoder_hidden_size	int	128	(undocumented)
decoder_layers	int	2	(undocumented)
futr_exog_size	int	0	(undocumented)
hist_exog_size	int	0	(undocumented)
stat_exog_size	int	0	(undocumented)
output_size	int	1	(undocumented)
recurrent	bool	False	(undocumented)
cell_type	str	"gru"	"gru" (default) or "elman". GRU gating gives meaningfully better accuracy; use "elman" only when you need exact weight parity with a vanilla RNN checkpoint.
layer_norm	bool	True	apply LayerNorm after each encoder layer (default True). Stabilises training especially with larger hidden sizes.

autoregressive_predict

model.autoregressive_predict

Multi-step autoregressive rollout for recurrent=True models.

The h-step rollout is implemented with jax.lax.scan so the entire loop is compiled as a single XLA op — significantly faster than the previous Python for loop which dispatched h separate model.apply calls.

Strategy: * History pass — feed the full history; capture encoder hidden state and last projected prediction. * Steps 1..h-1 — scan over the horizon, feeding the previous prediction (+ any futr_exog slice) one step at a time while threading the encoder hidden state.

Parameter	Type	Default	Description
model	RNN	-	:class:RNN instance with config.recurrent=True.
params	-	-	parameters from model.init or weight transfer.
insample_y	jnp.ndarray	-	[B, L, 1] historic targets.
h	int	-	number of forecast steps.
hist_exog	Optional[jnp.ndarray]	None	[B, L, X] historic exogenous, or None.
futr_exog	Optional[jnp.ndarray]	None	[B, L+H, F] exogenous covering history and horizon, or None.
stat_exog	Optional[jnp.ndarray]	None	[B, S] static exogenous, or None.

Returns: jnp.ndarray ([B, h, output_size] predictions.) Raises: ValueError (if model.config.recurrent is not True).

MLP

model.MLP · inherits flax.linen.Module

MLP head matching Nixtla NeuralForecast's _modules.MLP.

__init__(self, out_features, hidden_size, num_layers, dropout_rate=0.0)

Parameter	Type	Default	Description
out_features	int	-	(undocumented)
hidden_size	int	-	(undocumented)
num_layers	int	-	(undocumented)
dropout_rate	float	0.0	(undocumented)

ElmanRNNCell

model.ElmanRNNCell · inherits flax.linen.Module

One layer of an Elman RNN.

Computes h_t = act(W_ih x_t + b_ih + W_hh h_{t-1} + b_hh).

(carry, x) -> (new_carry, output) so the cell is directly compatible with :func:flax.linen.scan over the time axis.

__init__(self, hidden_size, activation='tanh', use_bias=True)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
activation	str	"tanh"	(undocumented)
use_bias	bool	True	(undocumented)

GRUCell

model.GRUCell · inherits flax.linen.Module

One layer of a Gated Recurrent Unit.

Same (carry, x) -> (new_carry, output) interface as :class:ElmanRNNCell — drop-in for :func:flax.linen.scan.

__init__(self, hidden_size, use_bias=True)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
use_bias	bool	True	(undocumented)

RNNEncoder

model.RNNEncoder · inherits flax.linen.Module

Stacked Elman/GRU encoder.

Each layer is unrolled with flax.linen.scan. Optional LayerNorm and dropout are applied between layers (not after the final layer, matching PyTorch convention).

__init__(self, hidden_size, num_layers, activation='tanh', use_bias=True, dropout_rate=0.0, cell_type='gru', layer_norm=True)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
num_layers	int	-	(undocumented)
activation	str	"tanh"	(undocumented)
use_bias	bool	True	(undocumented)
dropout_rate	float	0.0	(undocumented)
cell_type	str	"gru"	"gru" (default) or "elman".
layer_norm	bool	True	apply LayerNorm after each layer's output sequence.