sinusoid_position_embedding

chronax.informer_layers.sinusoid_position_embedding

Fixed (non-learned) sinusoidal position code, shape [L, hidden_size].

Standard transformer encoding: even columns are sin(pos*freq), odd columns are cos(pos*freq), with freq geometrically spaced via exp(-ln(10000)*i/hidden_size). L and hidden_size are static Python ints, so this constant-folds under jit; there is no learned parameter and no cached max-length buffer (unlike the torch reference's preallocated pe), matching this port's params-only nnx.Module style. hidden_size may be odd: 0::2/1::2 need equally-sized targets, so the sin/cos grids are built on a one-wider padded buffer when hidden_size is odd, then sliced back to hidden_size.

Parameter	Type	Default	Description
L	int	-	(undocumented)
hidden_size	int	-	(undocumented)

Returns: jnp.ndarray

TokenEmbedding

chronax.informer_layers.TokenEmbedding

Circular 1-D conv token embedding (NF TokenEmbedding).

A kernel-3 conv with padding="CIRCULAR" is length-preserving and wraps values at the sequence boundary, matching torch's Conv1d(kernel_size=3, padding=1, padding_mode='circular'). Flax is feature-last ([B, L, C]), so — unlike a literal torch port — no axis permutes are needed. No bias (NF default); Kaiming-normal init with fan_in = 3 * c_in.

__init__(self, c_in, hidden_size, *, rngs)

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

__call__(self, x)

x: [B, L, c_in] -> [B, L, hidden_size].

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

Returns: jnp.ndarray

TimeFeatureEmbedding

chronax.informer_layers.TimeFeatureEmbedding

Bias-free linear projection of time-feature marks (NF TimeFeatureEmbedding).

__init__(self, input_size, hidden_size, *, rngs)

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

__call__(self, x)

x: [B, L, input_size] -> [B, L, hidden_size].

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

Returns: jnp.ndarray

DataEmbedding

chronax.informer_layers.DataEmbedding

Token + sinusoidal-position (+ optional time-feature) embedding (NF DataEmbedding).

Sums the circular-conv token embedding, the fixed sinusoidal position code, and — only when exog_input_size > 0 — a bias-free linear embedding of future-known time-feature marks, then applies dropout. temporal_embedding is None (rather than a zero-input layer) when there are no exogenous marks, so x_mark may be passed as None in that case.

__init__(self, *, c_in, exog_input_size, hidden_size, dropout, rngs)

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

__call__(self, x, x_mark, deterministic)

x: [B, L, c_in], x_mark: [B, L, exog_input_size] or None -> [B, L, hidden_size].

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

Returns: jnp.ndarray

AttentionLayer

chronax.informer_layers.AttentionLayer

Multi-head ProbSparse attention (NF AttentionLayer + ProbAttention folded into one).

q/k/v/out projections are hidden_size -> hidden_size (torch init, bias=True — NF default), split into n_head heads of size hidden_size // n_head. Unlike itransformer's AttentionLayer (self-attention only), queries and keys/values here may have DIFFERENT lengths (L_Q vs L_K) — this is what lets the decoder's cross-attention run queries from the decoder against keys/values from the encoder. The inner attention is :func:_prob_attention rather than dense softmax attention.

Deliberate correctness deviation from Nixtla main: NF's AttentionLayer.forward flattens the ProbAttention context straight from its [B, H, L, E] layout (out = out.view(B, L, -1)), which silently interleaves the head and time axes whenever H > 1 (reinterpreting [B, H, L, E]-ordered memory as [B, L, H*E] mixes each head's features with the wrong time steps). The original Informer2020 code (and every other multi-head attention layer in this repo, e.g. itransformer_module.py's AttentionLayer) transposes to [B, L, H, E] before flattening. We follow that correct convention — ctx.transpose(0, 2, 1, 3).reshape(B, L_Q, hidden_size) — here rather than reproduce NF's bug.

No dropout (_prob_attention never applies its own), no deterministic arg.

__init__(self, *, hidden_size, n_head, factor, mask_flag, rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
n_head	int	-	(undocumented)
factor	int	-	(undocumented)
mask_flag	bool	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, queries, keys, values, *, sample_key)

queries: [B,L_Q,hid]; keys/values: [B,L_K,hid] -> [B,L_Q,hid].

Parameter	Type	Default	Description
queries	jnp.ndarray	-	(undocumented)
keys	jnp.ndarray	-	(undocumented)
values	jnp.ndarray	-	(undocumented)
sample_key	-	-	(undocumented)

Returns: jnp.ndarray

ConvLayer

chronax.informer_layers.ConvLayer

Self-attention distilling layer (NF ConvLayer).

Halves the sequence length between encoder stacks: a length-preserving circular conv (torch Conv1d(kernel_size=3, padding=1, padding_mode='circular') default init — same U(-1/sqrt(fan_in), 1/sqrt(fan_in)) bound as _torch_linear, with fan_in = 3 * c_in for both kernel and bias) feeds a BatchNorm + ELU, then a stride-2 max-pool. BatchNorm (rather than dropout/LayerNorm) is what the reference actually uses here, so — like patchtst_module.py's TSTEncoderLayer — the running-stat toggle is threaded as its own use_running_average flag; there is no dropout in this layer, so there is no deterministic argument.

The max-pool padding=((1, 1),) pads with -inf (flax's max_pool convention), matching torch's MaxPool1d(kernel_size=3, stride=2, padding=1) behavior on the boundary. Output length is (L-1)//2 + 1.

__init__(self, c_in, *, rngs)

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

__call__(self, x, *, use_running_average)

x: [B, L, C] -> [B, (L-1)//2 + 1, C].

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

Returns: jnp.ndarray

TransEncoderLayer

chronax.informer_layers.TransEncoderLayer

One Informer encoder layer (post-norm; NF EncoderLayer).

Adapted from itransformer_module.py's TransEncoderLayer (same post-norm skeleton: residual attention -> norm1 -> pointwise FFN -> residual -> norm2), swapping in the ProbSparse AttentionLayer (mask_flag=False, unmasked self-attention) and threading a sample_key for its data-independent key subsample. The two Conv1d(kernel_size=1) layers in NF are pointwise linear maps over the channel axis, so — as in itransformer's port — they are plain nnx.Linear (named conv1/conv2 to keep the NF correspondence visible), with FFN width conv_hidden_size.

__init__(self, *, hidden_size, n_head, conv_hidden_size, factor, dropout, activation='gelu', rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
n_head	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
factor	int	-	(undocumented)
dropout	float	-	(undocumented)
activation	str	"gelu"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, *, sample_key, deterministic)

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

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

Returns: jnp.ndarray

TransDecoderLayer

chronax.informer_layers.TransDecoderLayer

One Informer decoder layer (post-norm; NF DecoderLayer).

Same FFN/config shape as :class:TransEncoderLayer but with two attention sub-layers and three norms: masked ProbSparse self-attention (mask_flag=True) over the decoder's own sequence, then unmasked ProbSparse cross-attention (mask_flag=False) with queries from the decoder against keys/values from the encoder output (cross). Each attention call gets its own independent sampling key (self_key / cross_key) since the two AttentionLayers draw unrelated key subsamples.

__init__(self, *, hidden_size, n_head, conv_hidden_size, factor, dropout, activation='gelu', rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
n_head	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
factor	int	-	(undocumented)
dropout	float	-	(undocumented)
activation	str	"gelu"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x, cross, *, self_key, cross_key, deterministic)

x: [B, L_dec, hidden], cross: [B, L_enc, hidden] -> [B, L_dec, hidden].

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

Returns: jnp.ndarray