RevIN

itransformer_module.RevIN · inherits nnx.Module

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

Faithful to neuralforecast's iTransformer 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, *, subtract_last=False, affine=False, eps=1e-5, rngs)

(undocumented)

norm(self, x)

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

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].

Returns: jnp.ndarray

DataEmbeddingInverted

itransformer_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)

(undocumented)

__call__(self, x, deterministic)

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

Returns: jnp.ndarray

AttentionLayer

itransformer_module.AttentionLayer · inherits nnx.Module

Multi-head full (softmax dot-product) self-attention over the token axis.

Folds NF's AttentionLayer + FullAttention into one module: q/k/v/out projections (Linear(hidden -> n_heads*d_k), torch init) and a frozen scale d_k**-0.5 applied pre-softmax (NF scale = 1/sqrt(E)). Dropout is applied to the attention weights (NF FullAttention(attention_dropout=dropout)); the output projection carries no dropout (the residual dropout lives in the encoder layer). No causal mask (NF mask_flag=False), no residual-attention threading.

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

(undocumented)

__call__(self, x, deterministic)

x: [B, T, hidden] -> [B, T, hidden] (T = number of variate tokens).

Returns: (undocumented)

TransEncoderLayer

itransformer_module.TransEncoderLayer · inherits nnx.Module

One iTransformer encoder layer (post-norm).

NF TransEncoderLayer.forward::

new_x = attention(x, x, 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).

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

(undocumented)

__call__(self, x, deterministic)

(undocumented)

Returns: (undocumented)

TransEncoder

itransformer_module.TransEncoder · inherits nnx.Module

Stack of e_layers encoder layers plus a final LayerNorm (NF norm_layer).

__init__(self, *, e_layers, hidden_size, n_heads, d_ff, dropout, activation='gelu', rngs)

(undocumented)

__call__(self, x, deterministic)

(undocumented)

Returns: (undocumented)

ITransformerNet

itransformer_module.ITransformerNet · inherits nnx.Module

Full iTransformer 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, n_heads, e_layers, d_ff, dropout, use_norm, activation='gelu', rngs)

(undocumented)

__call__(self, x, deterministic)

(undocumented)

Returns: jnp.ndarray