xlstm_functions

chronax.xlstm_functions

High-level glue for the xLSTM / xLSTMTime forecaster.

Wires together the pure-functional pieces in xlstm_backend into a training entry point (xlstm_f) and a stateless decode wrapper (forecast_xlstm). Compiled XLA kernels are cached via functools.lru_cache keyed on the static XLSTMConfig plus hyperparameters that affect compilation, mirroring the cached-jit pattern used in ets/ets_functions.py and holt._get_holt_optimizer.

make_windows_ar(z, ctx_len, horizon_train)

AR-mode windows: input (T,), target shifted by 1. Both shape (n_win, ctx_len+horizon_train).

Parameter	Type	Default	Description
z	jnp.ndarray	-	(undocumented)
ctx_len	int	-	(undocumented)
horizon_train	int	-	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

make_windows_direct(z, ctx_len, horizon)

Direct-mode windows: input (ctx_len,), target (horizon,) starting at ctx_len.

Parameter	Type	Default	Description
z	jnp.ndarray	-	(undocumented)
ctx_len	int	-	(undocumented)
horizon	int	-	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

make_windows(z, ctx_len, horizon_train)

Backward-compat alias used by older call sites. Equivalent to make_windows_ar.

Parameter	Type	Default	Description
z	-	-	(undocumented)
ctx_len	-	-	(undocumented)
horizon_train	-	-	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

xlstm_f(z, cfg, key, *, n_epochs, batch_size, lr, weight_decay)

Train an xLSTM / xLSTMTime model on series z.

Series-level normalization (z-score) is the caller's responsibility when cfg.use_revin == False. When RevIN is on, pass the RAW series — the model handles normalization internally.

Parameter	Type	Default	Description
z	jnp.ndarray	-	Series to train on.
cfg	XLSTMConfig	-	(undocumented)
key	jax.Array	-	(undocumented)
n_epochs	int	-	(undocumented)
batch_size	int	-	(undocumented)
lr	float	-	(undocumented)
weight_decay	float	-	(undocumented)

Returns: dict ({"params", "losses", "cfg"}).

forecast_xlstm(model_dict, h)

Forecast h steps using stored trained params.

Returns a (h,) float32 array in the normalised scale when cfg.use_revin == False. The caller denormalises with the stored mu and std. When RevIN is on, returns the original-scale forecast (denormalisation happens inside the model).

Parameter	Type	Default	Description
model_dict	dict	-	(undocumented)
h	int	-	(undocumented)

Returns: jnp.ndarray