stemgnn_module.py

Flax NNX modules for the StemGNN forecaster (port of neuralforecast's StemGNN).

GLU

chronax.stemgnn_module.GLU · inherits nnx.Module

Gated linear unit: linear_left(x) * sigmoid(linear_right(x)).

The two linear maps are fused into one Linear(in, 2*out) whose output is split — identical math (concatenated output columns), same init law (both halves drawn from U(±1/sqrt(in))), and half the GEMM dispatches on the hot path.

__init__(self, in_features, out_features, *, rngs)

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

__call__(self, x) -> jnp.ndarray

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

TorchGRU

chronax.stemgnn_module.TorchGRU · inherits nnx.Module

Single-layer GRU with torch nn.GRU weight layout and gate math.

Parameters mirror torch exactly for pure-copy weight transplant: w_ih [3H, I], w_hh [3H, H], separate b_ih [3H]/b_hh [3H], gates packed r|z|n. The candidate gate keeps b_hn INSIDE the reset product — n = tanh(i_n + b_in + r * (h_n + b_hn)) — which is why flax's nnx.GRUCell (fused single bias) is NOT transplant-exact and is not used. All params init U(±1/sqrt(H)) (torch GRU default).

__init__(self, input_size, hidden_size, *, rngs)

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

__call__(self, x) -> jnp.ndarray

x: [seq, B, I] -> full output sequence [seq, B, H]; h0 = 0.

StemGNN runs the sequence over the NODE axis (seq = n_series), so this scan has length N — a single step in the univariate wrapper.

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

StockBlockLayer

chronax.stemgnn_module.StockBlockLayer · inherits nnx.Module

One StemGNN block: GFT -> Spe-Seq cell -> per-order graph-conv kernel -> sigmoid-gated forecast head (+ backcast head on block 0 only).

stack_cnt is the block index (0 or 1). Block 1 has no backcast Linear but still carries a backcast_short_cut that never influences its forward; it is built unconditionally to keep the parameter layout aligned with the reference for weight transplant.

__init__(self, time_step, unit, multi_layer, stack_cnt, *, rngs)

Parameter	Type	Default	Description
time_step	int	-	(undocumented)
unit	int	-	(undocumented)
multi_layer	int	-	(undocumented)
stack_cnt	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

spe_seq_cell(self, x) -> jnp.ndarray

[B, 4, N, L] -> [B, 4, N, S]: DFT over the ORDER axis, GLU stacks on real/imag, irfft back (top bin cropped).

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

__call__(self, x, mul_L)

x: [B, N, L]; mul_L: [4, N, N] -> (forecast [B, N, L], backcast [B, N, L] | None). The single input channel is dropped (in_channel == 1 always), so the einsums are the squeezed equivalents of the reference's broadcast matmuls.

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

StemGNNNet

chronax.stemgnn_module.StemGNNNet · inherits nnx.Module

Full StemGNN network forward.

insample_z [B, L, N] (scaled) -> [B, h, outputsize_multiplier * N]. At N>1 the final reshape interleaves quantile heads and series; the univariate wrapper always runs N=1, where it collapses to [B, h, mult].

__init__(self, *, h, input_size, n_series=1, n_stacks=2, multi_layer=5, dropout_rate=0.5, leaky_rate=0.2, outputsize_multiplier=1, chebyshev_first_term='nf_zero', rngs)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
n_series	int	1	(undocumented)
n_stacks	int	2	(undocumented)
multi_layer	int	5	(undocumented)
dropout_rate	float	0.5	(undocumented)
leaky_rate	float	0.2	(undocumented)
outputsize_multiplier	int	1	(undocumented)
chebyshev_first_term	str	'nf_zero'	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, insample_z, dropout_key=None, deterministic=True) -> jnp.ndarray

insample_z: [B, L, N] scaled -> [B, h, mult * N] scaled.

deterministic=True (inference/eval) needs no key — dropout is off. Training passes a per-step dropout key.

Parameter	Type	Default	Description
insample_z	jnp.ndarray	-	(undocumented)
dropout_key	-	None	(undocumented)
deterministic	bool	True	(undocumented)