RevIN

chronax.softs_module.RevIN · inherits nnx.Module

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

Faithful to neuralforecast's SOFTS use_norm block, which centers on the per-window MEAN (subtract_last=False), divides by sqrt(var + eps) with population variance (unbiased=False / ddof=0), and applies no learnable affine. Statistics are returned explicitly rather than cached, so the module is pure and vmap/scan-safe.

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

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

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

Normalizes the input x.

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array of shape [B, L, C].

Returns: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] (z: [B, L, C], loc: [B, 1, C], scale: [B, 1, C]).

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

Invert norm.

Parameters:

Parameter	Type	Default	Description
z	jnp.ndarray	-	Normalized input array.
loc	jnp.ndarray	-	Location statistics (mean or last value).
scale	jnp.ndarray	-	Scale statistics (standard deviation).

Returns: jnp.ndarray (The denormalized array).

DataEmbeddingInverted

chronax.softs_module.DataEmbeddingInverted · inherits nnx.Module

Inverted embedding: each variate's lookback becomes a token.

Mirrors NF DataEmbedding_inverted: permute [B, L, N] -> [B, N, L] then Linear(input_size -> hidden_size) (the lookback length is the feature dim), followed by dropout. Exogenous/time marks are not supported (univariate, no covariates), so the x_mark concat path is omitted.

__init__(self, *, input_size, hidden_size, dropout, rngs: nnx.Rngs)

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

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

Transforms input into tokens.

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array of shape [B, L, N].
deterministic	bool	-	Whether to run in deterministic (inference) mode (disables dropout).

Returns: jnp.ndarray (Tokens of shape [B, N, hidden_size]).

STAD

chronax.softs_module.STAD · inherits nnx.Module

STar Aggregate-Dispatch module — SOFTS's series-core fusion block.

Replaces the encoder's self-attention. Given per-series tokens input: [B, C, d_series] (d_series = hidden_size, C = number of variate tokens):

1. Set FFN h = gen2(gelu(gen1(input))) -> [B, C, d_core].
2. Aggregate into a core (pool across the C axis into a single representation, then broadcast it back to all C series): - train (deterministic=False): STOCHASTIC pooling — for each (batch, core-dim) sample one series index from softmax(h) over the C axis (NF torch.multinomial) and gather that series' value. - eval (deterministic=True): the softmax-weighted mean over the C axis (NF's inference branch).
3. Dispatch + fuse output = gen4(gelu(gen3([input, core]))) -> [B, C, d_series].

Cost is O(C) in the number of series, versus O(C^2) for attention — the point of SOFTS. The multinomial sample draws a fresh key from rngs each forward; under the training nnx.scan the key stream is threaded through the carry exactly like nnx.Dropout's, so successive steps sample independently.

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

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

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

Processes input tokens using the STAD mechanism.

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input tokens of shape [B, C, hidden].
deterministic	bool	-	If False, uses stochastic multinomial pooling (training); if True, uses deterministic softmax-weighted mean (inference).

Returns: jnp.ndarray (Output tokens of shape [B, C, hidden]).

TransEncoderLayer

chronax.softs_module.TransEncoderLayer · inherits nnx.Module

One SOFTS encoder layer (post-norm), STAD in place of attention.

The two Conv1d(kernel_size=1) layers are mathematically pointwise linear maps over the channel axis, so they are implemented as nnx.Linear (no transpose needed). activation is exact GELU (NF passes F.gelu).

__init__(self, *, hidden_size, d_core, d_ff, dropout, activation="gelu", rngs: nnx.Rngs)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
d_core	-	-	(undocumented)
d_ff	-	-	(undocumented)
dropout	-	-	(undocumented)
activation	str	"gelu"	Activation function name.
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, deterministic: bool)

Performs a forward pass through the encoder layer.

Parameters:

Parameter	Type	Default	Description
x	-	-	Input tensor.
deterministic	bool	-	Whether to run in deterministic mode (affects STAD pooling and dropout).

Returns: - (Output tensor).

TransEncoder

chronax.softs_module.TransEncoder · inherits nnx.Module

Stack of e_layers encoder layers. NO final LayerNorm — see below.

NF's common._modules.TransEncoder takes an OPTIONAL norm_layer and applies it only if self.norm is not None. SOFTS builds the encoder positionally: TransEncoder([TransEncoderLayer(STAD(...), ...) for l in range(e_layers)]) with no norm_layer argument, so the reference has no final normalization and the encoder output feeds projection directly — NF's state_dict has no encoder.norm.* entry at all.

The difference is subtler than it looks, and worth stating precisely: every layer already ENDS in norm2, so an extra final LayerNorm is near-identity at initialization (re-normalizing an already-normalized vector). What it is not is free — its scale and bias are learnable, so it would hand the port 2 * hidden_size trainable parameters the reference does not have, and a learned per-feature affine applied immediately before the projector.

__init__(self, *, e_layers, hidden_size, d_core, d_ff, dropout, activation="gelu", rngs: nnx.Rngs)

Parameter	Type	Default	Description
e_layers	-	-	(undocumented)
hidden_size	-	-	(undocumented)
d_core	-	-	(undocumented)
d_ff	-	-	(undocumented)
dropout	-	-	(undocumented)
activation	str	"gelu"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, deterministic: bool)

Performs a forward pass through all encoder layers.

Parameters:

Parameter	Type	Default	Description
x	-	-	Input tensor.
deterministic	bool	-	Whether to run in deterministic mode.

Returns: - (Output tensor).

SOFTSNet

chronax.softs_module.SOFTSNet · inherits nnx.Module

Full SOFTS backbone: (RevIN) -> invert-embed -> encoder -> project -> (denorm).

I/O mirrors the other Chronax neural nets: __call__(x: [B, L, 1]) -> [B, h, 1]. Written N-generically: an [B, L, N] input yields [B, h, N].

__init__(self, *, h, input_size, hidden_size, d_core, e_layers, d_ff, dropout, use_norm, activation="gelu", rngs: nnx.Rngs)

Parameter	Type	Default	Description
h	-	-	(undocumented)
input_size	-	-	(undocumented)
hidden_size	-	-	(undocumented)
d_core	-	-	(undocumented)
e_layers	-	-	(undocumented)
d_ff	-	-	(undocumented)
dropout	-	-	(undocumented)
use_norm	-	-	(undocumented)
activation	str	"gelu"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

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

Performs a forward pass through the SOFTS network.

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array of shape [B, L, N].
deterministic	bool	-	Whether to run in deterministic mode.

Returns: jnp.ndarray (Forecast output of shape [B, h, N]).