TokenEmbedding

timemixer_module.TokenEmbedding · inherits nnx.Module

k=3 circular-padded 1-D conv token embedding, computed as three rolled matmuls. Weight shape [hidden, c_in, 3] (torch conv layout), no bias, Kaiming-normal fan_in init. x: [B, T, c_in] -> [B, T, hidden].

__init__(self, c_in: int, hidden_size: int, *, rngs: nnx.Rngs)

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

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

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

RevIN

timemixer_module.RevIN · inherits nnx.Module

Reversible instance normalization with optional learnable affine.

Per-window mean over time (subtract_last unsupported here — the reference instantiates TimeMixer's copies with the default mean mode), population variance, sqrt(var + eps) scale. The affine inverse divides by gamma + eps**2 — eps squared, an inherited reference quirk. With non_norm=True both directions pass through unchanged while the affine parameters still exist (reference parity: it constructs them regardless). Statistics are returned explicitly so the module stays pure and vmap/scan-safe; they carry no parameter dependence, so no stop-gradient is needed for gradient equivalence with the reference's detach.

__init__(self, num_features: int, *, affine: bool = True, non_norm: bool = False, eps: float = 1e-5, rngs: nnx.Rngs)

Parameter	Type	Default	Description
num_features	int	-	(undocumented)
affine	bool	True	(undocumented)
non_norm	bool	False	(undocumented)
eps	float	1e-5	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

norm(self, x: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

x: [B, T, C] -> (z, loc [B,1,C], scale [B,1,C]).

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

denorm(self, z: jnp.ndarray, loc: jnp.ndarray, scale: jnp.ndarray) -> jnp.ndarray

Parameter	Type	Default	Description
z	jnp.ndarray	-	(undocumented)
loc	jnp.ndarray	-	(undocumented)
scale	jnp.ndarray	-	(undocumented)

MultiScaleSeasonMixing

timemixer_module.MultiScaleSeasonMixing · inherits nnx.Module

Bottom-up (fine -> coarse) season mixing across scales.

Inputs/outputs are lists over scales; tensors are [B, d, T_i] inside the chain and [B, T_i, d] in the returned list (reference layout).

__init__(self, widths: list[int], *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
widths	list[int]	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, season_list: list[jnp.ndarray]) -> list[jnp.ndarray]

Parameter	Type	Default	Description
season_list	list[jnp.ndarray]	-	(undocumented)

MultiScaleTrendMixing

timemixer_module.MultiScaleTrendMixing · inherits nnx.Module

Top-down (coarse -> fine) trend mixing across scales (reversed lists).

__init__(self, widths: list[int], *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
widths	list[int]	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, trend_list: list[jnp.ndarray]) -> list[jnp.ndarray]

Parameter	Type	Default	Description
trend_list	list[jnp.ndarray]	-	(undocumented)

PastDecomposableMixing

timemixer_module.PastDecomposableMixing · inherits nnx.Module

One PDM block: per-scale decomposition, optional channel-crossing MLP, season/trend mixing, and (channel-independent mode only) a residual through the output-crossing MLP. The reference also constructs a LayerNorm it never applies — kept here, unused, for parameter-inventory parity.

__init__(self, *, widths: list[int], d_model: int, d_ff: int, dropout: float, channel_independence: int, decomp_method: str, moving_avg: int, top_k: int, rngs: nnx.Rngs)

Parameter	Type	Default	Description
widths	list[int]	-	(undocumented)
d_model	int	-	(undocumented)
d_ff	int	-	(undocumented)
dropout	float	-	(undocumented)
channel_independence	int	-	(undocumented)
decomp_method	str	-	(undocumented)
moving_avg	int	-	(undocumented)
top_k	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x_list: list[jnp.ndarray]) -> list[jnp.ndarray]

Parameter	Type	Default	Description
x_list	list[jnp.ndarray]	-	(undocumented)

TimeMixerNet

timemixer_module.TimeMixerNet · inherits nnx.Module

Full TimeMixer: multi-scale downsampling -> per-scale RevIN -> (channel-dependent mode) season/trend pre-decomposition -> token embedding -> e_layers PDM blocks -> per-scale horizon predictors -> scale sum -> denormalization with the finest scale's statistics.

__call__(insample_y [B, L, N], deterministic) -> [B, h, N * mult] (the mult > 1 head is the reference's distr_output Linear over the series axis).

__init__(self, *, h: int, input_size: int, n_series: int, d_model: int = 32, d_ff: int = 32, dropout: float = 0.1, e_layers: int = 4, top_k: int = 5, decomp_method: str = "moving_avg", moving_avg: int = 25, channel_independence: int = 0, down_sampling_layers: int = 1, down_sampling_window: int = 2, down_sampling_method: str = "avg", use_norm: bool = True, outputsize_multiplier: int = 1, rngs: nnx.Rngs)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
n_series	int	-	(undocumented)
d_model	int	32	(undocumented)
d_ff	int	32	(undocumented)
dropout	float	0.1	(undocumented)
e_layers	int	4	(undocumented)
top_k	int	5	(undocumented)
decomp_method	str	"moving_avg"	(undocumented)
moving_avg	int	25	(undocumented)
channel_independence	int	0	(undocumented)
down_sampling_layers	int	1	(undocumented)
down_sampling_window	int	2	(undocumented)
down_sampling_method	str	"avg"	(undocumented)
use_norm	bool	True	(undocumented)
outputsize_multiplier	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, insample_y: jnp.ndarray, *, deterministic: bool = True) -> jnp.ndarray

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