PositionalEmbedding

softssharp_module.PositionalEmbedding · inherits nnx.Module

Additive sinusoidal encoding over the token (variate) axis.

__call__(x, scale) returns x + scale * table[:, :x.shape[1]], matching NF PositionalEmbedding.forward. scale is a scalar (array or float) so the caller can pass a learnable parameter, a gated parameter, or a constant.

__init__(self, d_series: int, max_len: int = 5000)

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

x: [B, C, d_series] -> same shape. Requires C <= max_len.

STADSharp

softssharp_module.STADSharp · inherits nnx.Module

STAD with stochastic variable-position encoding — the SOFTSSharp block.

Given per-series tokens input: [B, C, d_series] (d_series = hidden_size, C = number of variate tokens):

0. Variable-position encoding (new vs SOFTS): add the sinusoidal table indexed by variate position, scaled by the learnable pe_scale. Applied with probability pe_keep_prob in training (one Bernoulli per forward, shared across the batch); always applied at inference, scaled by pe_keep_prob * pe_scale. The encoded tensor is what feeds both the set FFN and the later fusion concat, exactly as in NF.
1. Set FFN h = gen2(dropout1(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). followed by dropout2 on the dispatched core (new vs SOFTS).
3. Dispatch + fuse output = gen4(dropout3(gelu(gen3([input, core])))) -> [B, C, d_series] (dropout3 new vs SOFTS).

Cost stays O(C) in the number of series, versus O(C^2) for attention. The multinomial sample and the position-encoding Bernoulli both draw fresh keys from rngs each forward; under the training nnx.scan the key streams are threaded through the carry exactly like nnx.Dropout's, so successive steps sample independently.

__init__(self, *, hidden_size, d_core, dropout: float = 0.1, pe_keep_prob: float = 0.5, pe_max_len: int = 5000, rngs: nnx.Rngs)

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

NF STADSharp.add_positional_embedding, as a branch-free JAX gate.

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

x: [B, C, hidden] -> [B, C, hidden].

RevIN

softssharp_module.RevIN · inherits nnx.Module

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

Faithful to neuralforecast's SOFTSSharp 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)

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

x: [B, L, C] -> (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. z: [B, T, C] with broadcastable loc/scale [B, 1, C].

DataEmbeddingInverted

softssharp_module.DataEmbeddingInverted · inherits nnx.Module

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

Mirrors NF DataEmbedding_inverted (SOFTSSharp imports SOFTS's verbatim): 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)

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

x: [B, L, N] -> tokens: [B, N, hidden_size].

TransEncoderLayer

softssharp_module.TransEncoderLayer · inherits nnx.Module

One SOFTSSharp encoder layer (post-norm), STADSharp in place of attention.

NF TransEncoderLayer.forward (with STADSharp as the "attention")::

new_x, _ = STADSharp(x)
x = x + dropout(new_x)
y = x = norm1(x)
y = dropout(activation(conv1(y)))   # conv1d kernel=1 == pointwise Linear
y = dropout(conv2(y))
return norm2(x + y)

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). Note the residual adds the layer input BEFORE the position encoding — the encoding lives entirely inside STADSharp, as in NF.

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

__call__(self, x, deterministic: bool)

TransEncoder

softssharp_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. SOFTSSharp (like SOFTS) builds the encoder positionally:

TransEncoder([TransEncoderLayer(STADSharp(...), ...) 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.

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. That is a real architectural divergence, so it is omitted.

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

__call__(self, x, deterministic: bool)

SOFTSSharpNet

softssharp_module.SOFTSSharpNet · inherits nnx.Module

Full SOFTSSharp 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], so a multivariate path can reuse it later; only the :class:SOFTSSharp wrapper fixes N = 1.

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

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

x: [B, L, N] -> [B, h, N].