chronax.models.tsmixer.forecaster

High-level fit/forecast adapter around the Chronax TSMixer.

TSMixer is a multivariate model — all N series must be provided together as a single [T, N] array (or a list of N equal-length 1-D arrays). A single shared model processes all channels simultaneously.

`TSMixerForecaster`

chronax.models.tsmixer.forecaster.TSMixerForecaster · inherits BaseForecaster

High-level fit/forecast wrapper around :class:chronax.models.tsmixer.TSMixer.

init(self, h: int, input_size: int = -1, n_series: int = 1, n_block: int = 2, ff_dim: int = 64, dropout: float = 0.1, revin: bool = True, revin_affine: bool = True, temporal_norm_momentum: float = 0.05, feature_norm_momentum: float = 0.05, use_batchnorm: bool = True, use_global_skip: bool = False, *, max_steps: int = 1000, learning_rate: float = 0.001, batch_size: int = 32, random_seed: int = 0, alias: str = 'TSMixer', loss: Union[str, LossFn] = 'mae', scale: bool = True, grad_clip: float = 0.0, use_lr_schedule: bool = False, weight_decay: float = 0.0, val_fraction: float = 0.1, val_check_steps: int = 100)

Parameter	Type	Default	Description
`h`	`int`	-	forecast horizon.
`input_size`	`int`	`-1`	history window length; -1 (default) uses `3 * h`.
`n_series`	`int`	`1`	number of time series (channels). Must match the number of channels in the data passed to :meth:`fit`.
`n_block`	`int`	`2`	forwarded to :class:`TSMixerConfig`.
`ff_dim`	`int`	`64`	forwarded to :class:`TSMixerConfig`.
`dropout`	`float`	`0.1`	forwarded to :class:`TSMixerConfig`.
`revin`	`bool`	`True`	forwarded to :class:`TSMixerConfig`.
`revin_affine`	`bool`	`True`	forwarded to :class:`TSMixerConfig`.
`temporal_norm_momentum`	`float`	`0.05`	forwarded to :class:`TSMixerConfig`.
`feature_norm_momentum`	`float`	`0.05`	forwarded to :class:`TSMixerConfig`.
`use_batchnorm`	`bool`	`True`	forwarded to :class:`TSMixerConfig`.
`use_global_skip`	`bool`	`False`	forwarded to :class:`TSMixerConfig`.
`max_steps`	`int`	`1000`	number of optimiser steps performed by :meth:`fit`.
`learning_rate`	`float`	`0.001`	Adam / AdamW peak learning rate.
`batch_size`	`int`	`32`	number of windows per training step.
`random_seed`	`int`	`0`	PRNG seed for parameter init and window sampling.
`alias`	`str`	`"TSMixer"`	display name for external reporting.
`loss`	`Union[str, LossFn]`	`"mae"`	registered name (`"mae"`, `"mse"`) from :mod:`chronax.models.tsmixer.loss` or a callable `(y, y_hat) -> scalar`.
`scale`	`bool`	`True`	if True, each training window is normalised by its own mean/std (per-window, not per-series global); inference context is normalised with the same convention.
`grad_clip`	`float`	`0.0`	global gradient-norm clipping threshold (0 = disabled).
`use_lr_schedule`	`bool`	`False`	if True, wrap Adam with warmup + cosine decay decaying to 1 % of `learning_rate` over `max_steps`.
`weight_decay`	`float`	`0.0`	if > 0, use AdamW instead of Adam.
`val_fraction`	`float`	`0.1`	(undocumented)
`val_check_steps`	`int`	`100`	(undocumented)

`fit(self, y: SeriesLike, *, verbose: bool = False) -> TSMixerForecaster`

Train the model on y.

Parameters:

Parameter	Type	Default	Description
`y`	`SeriesLike`	-	`[T, N]` multivariate array or list of N equal-length 1-D arrays. All channels are modelled jointly.
`verbose`	`bool`	`False`	print loss every `max(1, max_steps // 10)` steps.

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

`forecast(self, y: SeriesLike, h: Optional[int] = None) -> jnp.ndarray`

Produce h-step ahead forecasts.

Parameters:

Parameter	Type	Default	Description
`y`	`SeriesLike`	-	`[T, N]` series. The last `config.input_size` timesteps are used as the conditioning window.
`h`	`Optional[int]`	`None`	accepted for API symmetry; must equal `config.h` if provided.

Returns: jnp.ndarray ([h, N] forecast array, or [h] if the input was 1-D.)

`predict(self, h: Optional[int] = None, X: Optional[jnp.ndarray] = None, level: Optional[List[Union[int, float]]] = None) -> dict`

Forecast from the series passed to :meth:fit.

Satisfies the BaseForecaster contract: unlike :meth:forecast, this does not take y explicitly — it reuses the series cached at fit time.

Parameters:

Parameter	Type	Default	Description
`h`	`Optional[int]`	`None`	Forecast horizon. Defaults to `config.h` when None.
`X`	`Optional[jnp.ndarray]`	`None`	Reserved for future exogenous regressors; unused.
`level`	`Optional[List[Union[int, float]]]`	`None`	Not yet supported for this model.

Returns: dict ({"mean": jnp.ndarray}).

Raises: RuntimeError NotImplementedError

`fit_predict(self, y: SeriesLike, *, verbose: bool = False) -> jnp.ndarray`

Convenience: fit on y and immediately return forecasts.

Parameters:

Parameter	Type	Default	Description
`y`	`SeriesLike`	-	(undocumented)
`verbose`	`bool`	`False`	(undocumented)

Returns: jnp.ndarray.

`with_config(self, **overrides) -> TSMixerForecaster`

Return a fresh (unfitted) forecaster with overridden config fields.

Parameters:

Parameter	Type	Default	Description
`**overrides`	-	-	(undocumented)

Returns: TSMixerForecaster.

Properties

`fitted`

bool

`state`

TrainState