XLSTMConfig

chronax.xlstm_backend.XLSTMConfig

Static configuration for the xLSTM / xLSTMTime stack. Frozen so instances are hashable and usable as lru_cache keys and as static_argnames to jax.jit. Values flow into shape decisions and never get traced. New (xLSTMTime) fields default to the original mLSTM/AR recipe so the existing call sites and tests continue to produce bit-identical behavior.

__init__(self, embed_dim=64, num_heads=4, head_dim=16, num_layers=2, ctx_len=64, horizon_train=8, gate_clip=8.0, eps=1e-6, block_types=('mlstm', 'mlstm'), use_revin=False, revin_affine=True, use_decomposition=False, decomp_kernel=25, decode_mode='ar', horizon=8, use_conv1d_in_slstm=False, conv1d_kernel=4, slstm_forget_gate='exp', slstm_stabilizer='per_head')

Parameter	Type	Default	Description
embed_dim	int	64	D
num_heads	int	4	H
head_dim	int	16	Dh ; expected D == H * Dh
num_layers	int	2	L
ctx_len	int	64	C
horizon_train	int	8	HT (AR mode only)
gate_clip	float	8.0	(undocumented)
eps	float	1e-6	(undocumented)
block_types	Tuple[str, ...]	('mlstm', 'mlstm')	length must equal num_layers
use_revin	bool	False	(undocumented)
revin_affine	bool	True	(undocumented)
use_decomposition	bool	False	(undocumented)
decomp_kernel	int	25	(undocumented)
decode_mode	str	'ar'	"ar" | "direct"
horizon	int	8	static; required when decode_mode == "direct"
use_conv1d_in_slstm	bool	False	(undocumented)
conv1d_kernel	int	4	(undocumented)
slstm_forget_gate	str	'exp'	"exp" (current) | "sigmoid" (log-sigmoid forget)
slstm_stabilizer	str	'per_head'	"per_head" (current) | "per_cell" (per Eq.15)

BlockState

chronax.xlstm_backend.BlockState

Per-block mLSTM state.

Attributes

Attribute	Type	Description
C	jnp.ndarray	covariance matrix per head, shape (H, Dh, Dh). bf16 in compute.
n	jnp.ndarray	normalizer vector per head, shape (H, Dh). bf16.
m	jnp.ndarray	log-space stabilizer per head, shape (H,). float32.

SLSTMBlockState

chronax.xlstm_backend.SLSTMBlockState

Per-block sLSTM state.

Attributes

Attribute	Type	Description
h	jnp.ndarray	hidden state per head (needed for memory mixing via R), shape (H, Dh). bf16.
c	jnp.ndarray	scalar cell per unit, shape (H, Dh). bf16.
n	jnp.ndarray	normalizer, shape (H, Dh). bf16.
m	jnp.ndarray	log-space stabilizer, float32. Shape (H,) for slstm_stabilizer="per_head" (default) or (H, Dh) for "per_cell" (paper Eq.15).

init_block_state

chronax.xlstm_backend.init_block_state(cfg, dtype=jnp.bfloat16)

Parameter	Type	Default	Description
cfg	XLSTMConfig	-	(undocumented)
dtype	Any	jnp.bfloat16	(undocumented)

Returns: BlockState

init_slstm_block_state

chronax.xlstm_backend.init_slstm_block_state(cfg, dtype=jnp.bfloat16)

Parameter	Type	Default	Description
cfg	XLSTMConfig	-	(undocumented)
dtype	Any	jnp.bfloat16	(undocumented)

Returns: SLSTMBlockState

init_params

chronax.xlstm_backend.init_params(key, cfg)

Initialize the full parameter pytree in float32. Conditionally adds RevIN, decomposition kernel, and the direct head based on the XLSTMConfig flags. AR-mode out_proj (D,1) is always present so AR decode keeps working.

Parameter	Type	Default	Description
key	jax.Array	-	(undocumented)
cfg	XLSTMConfig	-	(undocumented)

Returns: dict

revin_normalize

chronax.xlstm_backend.revin_normalize(x, p, eps=1e-5)

Per-instance normalization. x: (T,) or (T, D_in).

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
p	dict	-	If p is a non-empty dict, applies learnable affine gamma * x_norm + beta.
eps	float	1e-5	(undocumented)

Returns: tuple Returns (x_norm, stats) where stats = {"mu": (1,) or (1, D_in), "sigma": same}.

revin_denormalize

chronax.xlstm_backend.revin_denormalize(y, p, stats)

Inverse of :func:revin_normalize. y: (H,) or (H, D_in).

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
p	dict	-	(undocumented)
stats	dict	-	(undocumented)

Returns: jnp.ndarray

series_decompose

chronax.xlstm_backend.series_decompose(x, p, kernel)

Split x into trend and seasonal components via a learnable moving-average kernel. Edge-replication padding to preserve length.

Parameter	Type	Default	Description
x	jnp.ndarray	-	(T,) float.
p	dict	-	(undocumented)
kernel	int	-	(undocumented)

Returns: tuple Returns (trend: (T,), seasonal: (T,)).

block_init_state

chronax.xlstm_backend.block_init_state(cfg, layer_idx, dtype=jnp.bfloat16)

Parameter	Type	Default	Description
cfg	XLSTMConfig	-	(undocumented)
layer_idx	int	-	(undocumented)
dtype	Any	jnp.bfloat16	(undocumented)

Returns: BlockState or SLSTMBlockState

xlstm_forward

chronax.xlstm_backend.xlstm_forward(params, x_seq, cfg)

Full xLSTM / xLSTMTime forward pass.

Parameter	Type	Default	Description
params	Any	-	pytree of float32 arrays (master weights).
x_seq	jnp.ndarray	-	(T,) float32 or bf16. Univariate scalar series.
cfg	XLSTMConfig	-	(undocumented)

Returns: tuple Returns (preds, final_states). AR mode: preds shape (T,) float32 — per-position next-step prediction. Direct mode: preds shape (cfg.horizon,) float32 — direct multi-step forecast. final_states: tuple of per-block carry states (only meaningful in AR mode).

decode

chronax.xlstm_backend.decode(params, z_tail, h_steps, cfg)

Encode z_tail then autoregress h_steps forward (AR mode only).

Parameter	Type	Default	Description
params	Any	-	float32 pytree.
z_tail	jnp.ndarray	-	(ctx_len,) float32, normalized.
h_steps	int	-	Python int (static). Number of forecast steps.
cfg	XLSTMConfig	-	(must have decode_mode == "ar").

Returns: jnp.ndarray Returns preds: (h_steps,) float32, normalized predictions.

decode_direct

chronax.xlstm_backend.decode_direct(params, z_tail, cfg)

Direct multi-step forecast. z_tail: (ctx_len,) float32, possibly NOT pre-normalized when RevIN is enabled (RevIN runs inside xlstm_forward).

Parameter	Type	Default	Description
params	Any	-	(undocumented)
z_tail	jnp.ndarray	-	(undocumented)
cfg	XLSTMConfig	-	(undocumented)

Returns: jnp.ndarray Returns (cfg.horizon,) float32.