RevIN

patchtst_module.RevIN · inherits nnx.Module

Reversible instance normalization (Kim et al. 2022), per-window.

Faithful to neuralforecast's RevIN at PatchTST's defaults: centers on the last timestep (subtract_last=True), divides by sqrt(var + eps) with population variance (ddof=0), and applies no learnable affine (affine=False). Statistics are returned explicitly rather than cached, so the module is pure and vmap/scan-safe.

__init__(self, num_features, *, subtract_last=True, affine=False, eps=1e-5, rngs)

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

norm(self, x)

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

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

Returns: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

denorm(self, z, loc, scale)

Invert norm. z: [B, T, C] with broadcastable loc/scale [B, 1, C].

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

Returns: jnp.ndarray

compute_patch_num

patchtst_module.compute_patch_num

Number of patches after end-padding by stride (NF padding_patch='end').

Mirrors NF's int((input_size - patch_len) / stride + 1) + 1 exactly (truncation toward zero, not Python floor) so it stays correct even if called with an unclamped patch_len; callers normally pass the clamped patch_len = min(input_size + stride, patch_len).

Parameter	Type	Default	Description
input_size	int	-	(undocumented)
patch_len	int	-	(undocumented)
stride	int	-	(undocumented)

Returns: int

patchify

patchtst_module.patchify

x: [B, L] -> patches: [B, patch_num, patch_len].

Replicates torch ReplicationPad1d((0, stride)) then unfold(dim=-1, size=patch_len, step=stride).

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

Returns: jnp.ndarray

PatchEmbedding

patchtst_module.PatchEmbedding · inherits nnx.Module

Linear patch projection + learnable positional encoding + residual dropout.

Mirrors NF's W_P (Linear(patch_len -> hidden_size)) plus positional_encoding(pe='zeros', learn_pe=True) — a learnable [patch_num, hidden_size] parameter initialized Uniform(-0.02, 0.02). Receives pre-cut patches; patchify lives in PatchTSTNet.

__init__(self, *, patch_len, hidden_size, patch_num, dropout, rngs)

Parameter	Type	Default	Description
patch_len	-	-	(undocumented)
hidden_size	-	-	(undocumented)
patch_num	-	-	(undocumented)
dropout	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, patches, deterministic)

patches: [B, patch_num, patch_len] -> tokens: [B, patch_num, hidden_size].

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

Returns: jnp.ndarray

MultiHeadAttention

patchtst_module.MultiHeadAttention · inherits nnx.Module

Multi-head self-attention with Realformer residual-attention threading.

Faithful to NF's res_attention=True path: pre-softmax scores are added across layers via prev and the scaling factor d_k**-0.5 is a frozen constant (NF uses lsa=False). Hand-rolled (not jax.nn.dot_product_attention) because that fused path only matches NF's res_attention=False branch.

__init__(self, *, hidden_size, n_heads, attn_dropout, proj_dropout, rngs)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
n_heads	-	-	(undocumented)
attn_dropout	-	-	(undocumented)
proj_dropout	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, prev, deterministic)

x: [B, T, hidden] -> (out: [B, T, hidden], scores: [B, H, T, T]).

Parameter	Type	Default	Description
x	-	-	(undocumented)
prev	-	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: (undocumented)

TSTEncoderLayer

patchtst_module.TSTEncoderLayer · inherits nnx.Module

One transformer encoder layer: residual MHA + BatchNorm, then FFN + BatchNorm.

Post-norm (NF default pre_norm=False). BatchNorm with axis=-1 over [B, patch_num, hidden] matches NF's transpose/BatchNorm1d sandwich; momentum=0.9 matches torch BatchNorm1d's momentum=0.1 running-stat decay. (flax updates running-var with the biased/ddof=0 batch variance vs torch's unbiased/ddof=1 — negligible at the benchmark's effective batch size of thousands, and irrelevant to the parity gate, which loads NF's stats and only runs eval.)

__init__(self, *, hidden_size, n_heads, linear_hidden_size, dropout, attn_dropout, activation='gelu', rngs)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
n_heads	-	-	(undocumented)
linear_hidden_size	-	-	(undocumented)
dropout	-	-	(undocumented)
attn_dropout	-	-	(undocumented)
activation	str	'gelu'	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, prev, deterministic, use_running_average)

Parameter	Type	Default	Description
x	-	-	(undocumented)
prev	-	-	(undocumented)
deterministic	bool	-	(undocumented)
use_running_average	bool	-	(undocumented)

Returns: (undocumented)

TSTEncoder

patchtst_module.TSTEncoder · inherits nnx.Module

Stack of n_layers encoder layers, threading residual attention scores.

__init__(self, *, n_layers, hidden_size, n_heads, linear_hidden_size, dropout, attn_dropout, activation='gelu', rngs)

Parameter	Type	Default	Description
n_layers	-	-	(undocumented)
hidden_size	-	-	(undocumented)
n_heads	-	-	(undocumented)
linear_hidden_size	-	-	(undocumented)
dropout	-	-	(undocumented)
attn_dropout	-	-	(undocumented)
activation	str	'gelu'	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, deterministic, use_running_average)

Parameter	Type	Default	Description
x	-	-	(undocumented)
deterministic	bool	-	(undocumented)
use_running_average	bool	-	(undocumented)

Returns: (undocumented)

FlattenHead

patchtst_module.FlattenHead · inherits nnx.Module

Flatten and project to the horizon (univariate).

NF flattens its encoder output [B, hidden, patch_num] with nn.Flatten(start_dim=-2) — HIDDEN-major (hidden is the slow axis). Our encoder produces [B, patch_num, hidden], so we transpose to [B, hidden, patch_num] BEFORE flattening, otherwise the linear-head weight columns are permuted relative to NF and the parity gate fails (measured: ~1.8 patch-major vs <1e-4 hidden-major — the gate passes only hidden-major).

__init__(self, *, hidden_size, patch_num, h, head_dropout, rngs)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
patch_num	-	-	(undocumented)
h	-	-	(undocumented)
head_dropout	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, deterministic)

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

Returns: jnp.ndarray

PatchTSTNet

patchtst_module.PatchTSTNet · inherits nnx.Module

Full PatchTST backbone: RevIN -> patchify -> embed -> encoder -> head -> denorm.

I/O mirrors the GRU network: __call__(x: [B, L, 1]) -> [B, h, 1].

__init__(self, *, h, input_size, patch_len, stride, hidden_size, n_heads, encoder_layers, linear_hidden_size, dropout, fc_dropout, head_dropout, attn_dropout, revin, revin_affine, revin_subtract_last, activation='gelu', rngs)

Parameter	Type	Default	Description
h	-	-	(undocumented)
input_size	-	-	(undocumented)
patch_len	-	-	(undocumented)
stride	-	-	(undocumented)
hidden_size	-	-	(undocumented)
n_heads	-	-	(undocumented)
encoder_layers	-	-	(undocumented)
linear_hidden_size	-	-	(undocumented)
dropout	-	-	(undocumented)
fc_dropout	-	-	fc_dropout is accepted for NF-signature parity but inert: NF uses it only in the disabled pretrain head; the active residual dropout is dropout.
head_dropout	-	-	(undocumented)
attn_dropout	-	-	(undocumented)
revin	-	-	(undocumented)
revin_affine	-	-	(undocumented)
revin_subtract_last	-	-	(undocumented)
activation	str	'gelu'	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, deterministic, use_running_average)

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

Returns: jnp.ndarray