distilled_length

chronax.models.informer.informer_module.distilled_length

Sequence length remaining after n_conv distilling ConvLayers.

Each conv applies L -> (L-1)//2 + 1 (the stride-2 maxpool arithmetic from ConvLayer). Pure Python int arithmetic on static shapes, so callers (e.g. the training/model wrappers) can size buffers without invoking the net.

distilled_length(input_size, n_conv)

Parameter	Type	Default	Description
input_size	int	-	(undocumented)
n_conv	int	-	(undocumented)

Returns: int

TransEncoder

chronax.models.informer.informer_module.TransEncoder · inherits nnx.Module

Stack of encoder_layers ProbSparse attention layers, optionally distilled.

NF interleave (distil=True): attn -> conv -> attn -> conv -> ... -> attn -- encoder_layers attention layers alternating with encoder_layers - 1 ConvLayers, where the LAST attention layer has no conv after it. Without distilling, it is a plain attention stack (no length reduction). Either way a final LayerNorm closes the encoder (NF norm_layer). Each attention layer consumes its own entry of sample_keys (ProbSparse's data-independent key subsample, one per attention site); the ConvLayers consume the shared use_running_average flag (BatchNorm running-stat toggle).

__init__(self, *, encoder_layers, hidden_size, n_head, conv_hidden_size, factor, dropout, activation="gelu", distil, rngs)

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

__call__(self, x, *, sample_keys, deterministic, use_running_average)

x: [B, L, hidden] -> [B, L', hidden] (L' == L iff not distilled).

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

Returns: jnp.ndarray

TransDecoder

chronax.models.informer.informer_module.TransDecoder · inherits nnx.Module

Stack of decoder_layers masked-self + cross ProbSparse attention layers, plus head.

Each layer consumes a (self_key, cross_key) pair taken as consecutive entries of sample_keys, in order. After the stack, a final LayerNorm (NF norm_layer) precedes the output projection -- unlike TFTNet/ITransformerNet, the output head lives INSIDE the decoder (NF Decoder.projection convention), not as a separate top-level adapter.

__init__(self, *, decoder_layers, hidden_size, n_head, conv_hidden_size, factor, dropout, activation="gelu", c_out, rngs)

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

__call__(self, x, cross, *, sample_keys, deterministic)

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

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

Returns: jnp.ndarray

InformerNet

chronax.models.informer.informer_module.InformerNet · inherits nnx.Module

Full Informer backbone: embed -> distilling encoder -> generative decoder -> head.

I/O mirrors the other Chronax neural nets: __call__(insample_y: [B, L, 1], futr_exog: [B, L+h, F] | None) -> [B, h, outputsize_multiplier]. Encoder and decoder get SEPARATE DataEmbeddings (NF does not share embedding weights between them), both fed the same future-known time-feature marks sliced to their own window. The decoder input is NF's generative "start token" trick: the last label_len steps of the raw history concatenated with h zero placeholders, so a single forward pass produces all h horizon steps (no autoregressive loop). One sample_key is split into n_attn_sites = encoder_layers + 2 * decoder_layers independent keys per call -- one per ProbSparse attention site (encoder self-attention layers, in order, then decoder (self, cross) pairs, in order) -- so no key is ever reused across sites.

__init__(self, *, h, input_size, label_len, hidden_size=128, n_head=4, factor=3, conv_hidden_size=32, encoder_layers=2, decoder_layers=1, distil=True, dropout=0.05, activation="gelu", futr_exog_size=0, outputsize_multiplier=1, rngs)

Parameter	Type	Default	Description
h	-	-	(undocumented)
input_size	-	-	(undocumented)
label_len	-	-	(undocumented)
hidden_size	-	128	(undocumented)
n_head	-	4	(undocumented)
factor	-	3	(undocumented)
conv_hidden_size	-	32	(undocumented)
encoder_layers	-	2	(undocumented)
decoder_layers	-	1	(undocumented)
distil	-	True	(undocumented)
dropout	-	0.05	(undocumented)
activation	str	"gelu"	(undocumented)
futr_exog_size	-	0	(undocumented)
outputsize_multiplier	-	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, insample_y, futr_exog=None, *, sample_key, deterministic, use_running_average)

Parameter	Type	Default	Description
insample_y	-	-	(undocumented)
futr_exog	-	None	(undocumented)
sample_key	-	-	(undocumented)
deterministic	bool	-	(undocumented)
use_running_average	bool	-	(undocumented)

Returns: jnp.ndarray (Shape [B, h, c_out])