`xlstm_model.py`

User-facing :class:XLSTM forecaster.

`XLSTM`

chronax.models.xlstm_model.XLSTM · inherits BaseForecaster

xLSTM / xLSTMTime forecaster.

Attributes: uses_exog: False

`init(self, ctx_len=64, horizon_train=8, num_layers=2, embed_dim=64, num_heads=4, n_epochs=100, batch_size=32, lr=0.001, weight_decay=1e-05, gate_clip=8.0, seed=0, alias='XLSTM', conformal_params=None, block_types=None, use_revin=False, revin_affine=True, use_decomposition=False, decomp_kernel=25, decode_mode='ar', horizon=None, use_conv1d_in_slstm=False, conv1d_kernel=4, slstm_forget_gate='exp', slstm_stabilizer='per_head')`

Parameter	Type	Default	Description
`ctx_len`	`int`	`64`	Sliding-window context length and training horizon (AR mode).
`horizon_train`	`int`	`8`	Sliding-window context length and training horizon (AR mode).
`num_layers`	`int`	`2`	Stack depth and per-block width.
`embed_dim`	`int`	`64`	Stack depth and per-block width. `embed_dim` must be divisible by `num_heads`.
`num_heads`	`int`	`4`	Stack depth and per-block width.
`n_epochs`	`int`	`100`	training hyperparameters.
`batch_size`	`int`	`32`	training hyperparameters.
`lr`	`float`	`1e-3`	training hyperparameters.
`weight_decay`	`float`	`1e-5`	training hyperparameters.
`gate_clip`	`float`	`8.0`	Pre-gate clipping bound on input/forget gate logits.
`seed`	`int`	`0`	PRNG seed for parameter initialisation and batch sampling.
`alias`	`str`	`"XLSTM"`	Identifier used by :meth:`__repr__`.
`conformal_params`	`Optional[ConformalIntervals]`	`None`	If provided, conformity scores are computed and cached at fit time.
`block_types`	`Optional[Tuple[str, ...]]`	`None`	Per-layer block type; `None` (default) -> `("mlstm",) * num_layers`. Each entry must be `"mlstm"` or `"slstm"`.
`use_revin`	`bool`	`False`	Reversible instance normalization. When True, the outer z-score in `fit` is skipped (RevIN handles per-instance stats internally).
`revin_affine`	`bool`	`True`	Whether RevIN has learnable affine `gamma`/`beta`.
`use_decomposition`	`bool`	`False`	Learnable moving-average trend + seasonal split, processed through shared-weight branches and summed.
`decomp_kernel`	`int`	`25`	Moving-average kernel size. Paper default 25.
`decode_mode`	`str`	`"ar"`	`"ar"` -> autoregressive decode (per-position next-step head). `"direct"` -> direct linear head `D -> horizon` (xLSTMTime canonical).
`horizon`	`Optional[int]`	`None`	Required when `decode_mode == "direct"`. Caller's forecast `h` must equal this value.
`use_conv1d_in_slstm`	`bool`	`False`	Apply causal Conv1D before sLSTM recurrence.
`conv1d_kernel`	`int`	`4`	Causal-Conv1D kernel size.
`slstm_forget_gate`	`str`	`"exp"`	sLSTM forget-gate form. `"exp"` (default) preserves the original behavior; `"sigmoid"` uses the log-sigmoid forget gate of the xLSTM paper (Eq. 13-15, arXiv 2412.07752). sLSTM blocks only.
`slstm_stabilizer`	`str`	`"per_head"`	sLSTM log-space stabilizer granularity. `"per_head"` (default) keeps the original collapsed stabilizer; `"per_cell"` matches the paper (Eq. 15).

`fit(self, y, X=None) -> Self`

Builds the model configuration, handles normalization (or RevIN setup), trains the model, and optionally computes conformity scores if conformal_params were set.

Parameters:

Parameter	Type	Default	Description
`y`	`jnp.ndarray`	-	(undocumented)
`X`	`jnp.ndarray \\| None`	`None`	(undocumented)

Returns: Self (the fitted forecaster; sets self.model_).

`predict(self, h, X=None, level=None) -> dict`

Generates an $h$-step forecast using the fitted model parameters and the stored context (self.model_["z_tail"]).

Parameters:

Parameter	Type	Default	Description
`h`	`int`	-	(undocumented)
`X`	`jnp.ndarray \\| None`	`None`	(undocumented)
`level`	`list[int \\| float] \\| None`	`None`	(undocumented)

Returns: dict (The prediction results). Keys include {"mean": jnp.ndarray} (plus interval keys if level is set).

`forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> dict`

h-step forecast.

Fast path: when self.model_ is populated (e.g. inside conformal walk-forward), reuses trained params + decodes from the last ctx_len of y. Slow path (no prior fit): fits a temporary model on y then predicts.

Parameters:

Parameter	Type	Default	Description
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`X`	`jnp.ndarray \\| None`	`None`	(undocumented)
`X_future`	`jnp.ndarray \\| None`	`None`	(undocumented)
`level`	`list[int \\| float] \\| None`	`None`	(undocumented)
`fitted`	`bool`	`False`	(undocumented)

Returns: dict (The prediction results). Keys include {"mean": jnp.ndarray} (plus interval keys if level is set in slow path).