chronax.models.tsmixer

Chronax TSMixer — JAX/Flax port of Nixtla NeuralForecast's TSMixer model.

TSMixerConfig

chronax.models.tsmixer.TSMixerConfig

Configuration for the TSMixer model.

__init__(self, context_length=32, prediction_length=8, n_block=2, d_model=128, dropout_rate=0.1, activation=<function gelu at 0x...>)

Parameter	Type	Default	Description
context_length	int	32	(undocumented)
prediction_length	int	8	(undocumented)
n_block	int	2	(undocumented)
d_model	int	128	(undocumented)
dropout_rate	float	0.1	(undocumented)
activation	Callable	nn.gelu	(undocumented)

TemporalMixing

chronax.models.tsmixer.TemporalMixing · inherits flax.linen.Module

Temporal Mixing block (MLP applied across the time dimension).

__call__(self, x)

(undocumented)

Parameter	Type	Default	Description
x	Any	-	Input array of shape (batch_size, context_length, d_model).

Returns: Output array of shape (batch_size, context_length, d_model).

FeatureMixing

chronax.models.tsmixer.FeatureMixing · inherits flax.linen.Module

Feature Mixing block (MLP applied across the feature dimension).

__call__(self, x)

(undocumented)

Parameter	Type	Default	Description
x	Any	-	Input array of shape (batch_size, context_length, d_model).

Returns: Output array of shape (batch_size, context_length, d_model).

MixingLayer

chronax.models.tsmixer.MixingLayer · inherits flax.linen.Module

A single TSMixer block consisting of Temporal Mixing and Feature Mixing.

__call__(self, x, training)

Applies the mixing layer.

Parameter	Type	Default	Description
x	Any	-	Input array of shape (batch_size, context_length, d_model).
training	bool	-	Whether in training mode (for dropout).

Returns: Output array of shape (batch_size, context_length, d_model).

TSMixer

chronax.models.tsmixer.TSMixer · inherits flax.linen.Module

The main TSMixer model architecture.

__call__(self, x, training)

Forward pass of the TSMixer model.

Parameter	Type	Default	Description
x	Any	-	Input array of shape (batch_size, context_length, n_features).
training	bool	-	Whether in training mode (for dropout).

Returns: Output array of shape (batch_size, prediction_length, 1).

TSMixerForecaster

chronax.models.tsmixer.TSMixerForecaster · inherits BaseForecaster

TSMixer Forecaster implementation using JAX/Flax.

This forecaster wraps the TSMixer model and handles data preparation, training loops, and prediction according to the BaseForecaster contract.

Attributes: * uses_exog: False * alias: "tsmixer"

__init__(self, config=TSMixerConfig(), optimizer=optax.adam(1e-3), n_epochs=10, batch_size=32, seed=42)

Initializes the TSMixer Forecaster.

Parameter	Type	Default	Description
config	TSMixerConfig	TSMixerConfig()	Configuration object for the TSMixer model architecture.
optimizer	optax.GradientTransformation	optax.adam(1e-3)	Optax optimizer definition.
n_epochs	int	10	Number of training epochs.
batch_size	int	32	Size of the training batches.
seed	int	42	Random seed for initialization and training.

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

Trains the TSMixer model.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Time series data (history) of shape (n_time, n_series).
X	Optional[jnp.ndarray]	None	Optional exogenous variables (ignored as uses_exog=False).

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

predict(self, h, X=None, level=None) -> Dict[str, jnp.ndarray]

Generates point forecasts for the next h steps based on the fitted model.

Parameter	Type	Default	Description
h	int	-	Forecast horizon (must match config.prediction_length).
X	Optional[jnp.ndarray]	None	Optional exogenous variables (ignored).
level	Optional[Sequence[int]]	None	Prediction intervals (not supported, ignored).

Returns: A dictionary containing the 'mean' forecast. Example: {"mean": jnp.ndarray(shape=(h, n_series))}. Raises: RuntimeError, ValueError, NotImplementedError

forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> Dict[str, jnp.ndarray]

Generates forecasts using the provided history y.

Parameter	Type	Default	Description
y	jnp.ndarray	-	History time series data (n_time, n_series).
h	int	-	Forecast horizon (must match config.prediction_length).
X	Optional[jnp.ndarray]	None	Optional exogenous variables (ignored).
X_future	Optional[jnp.ndarray]	None	Optional future exogenous variables (ignored).
level	Optional[Sequence[int]]	None	Prediction intervals (not supported, ignored).
fitted	bool	False	If True, assumes the model is already fitted (ignored here).

Returns: A dictionary containing the 'mean' forecast. Example: {"mean": jnp.ndarray(shape=(h, n_series))}. Raises: RuntimeError, ValueError

TrainState

chronax.models.tsmixer.TrainState · inherits flax.training.train_state.TrainState

Custom TrainState to hold batch statistics (if needed) and apply_fn.

create_train_state(model, rng_key, dummy_input, optimizer) -> TrainState

chronax.models.tsmixer.create_train_state

Initializes the TrainState for the TSMixer model.

Parameter	Type	Default	Description
model	Any	-	(undocumented)
rng_key	Any	-	(undocumented)
dummy_input	Any	-	(undocumented)
optimizer	Any	-	(undocumented)

Returns: TrainState

train_step(state, x, y) -> Tuple[TrainState, Dict[str, jnp.ndarray]]

chronax.models.tsmixer.train_step

Performs a single training step (forward pass, loss calculation, gradient update).

Parameter	Type	Default	Description
state	TrainState	-	Current training state.
x	jnp.ndarray	-	Input context data (B, L, D).
y	jnp.ndarray	-	Target prediction data (B, P, D).

Returns: Tuple[TrainState, Dict[str, jnp.ndarray]] (updated_state, metrics_dict).

eval_step(state, x, y) -> Dict[str, jnp.ndarray]

chronax.models.tsmixer.eval_step

Performs a single evaluation step.

Parameter	Type	Default	Description
state	TrainState	-	Current training state.
x	jnp.ndarray	-	Input context data (B, L, D).
y	jnp.ndarray	-	Target prediction data (B, P, D).

Returns: Metrics dictionary.

scan_train_loop()

chronax.models.tsmixer.scan_train_loop

(undocumented helper function)

train_loop()

chronax.models.tsmixer.train_loop

(undocumented helper function)

create_windows(y, context_length, prediction_length) -> jnp.ndarray

chronax.models.tsmixer.create_windows

Creates sliding windows from the time series data for training.

The windows have shape (context_length + prediction_length, n_series).

Parameter	Type	Default	Description
y	jnp.ndarray	-	Time series data (n_time, n_series).
context_length	int	-	Length of the input context (L).
prediction_length	int	-	Length of the output prediction (P).

Returns: jnp.ndarray (Array of windows (n_windows, L+P, n_series)). Raises: ValueError

make_batch(windows, context_length) -> Tuple[jnp.ndarray, jnp.ndarray]

chronax.models.tsmixer.make_batch

Splits windows into input (X) and target (Y) arrays.

Parameter	Type	Default	Description
windows	jnp.ndarray	-	Array of windows (B, L+P, D).
context_length	int	-	Length of the input context (L).

Returns: Tuple[jnp.ndarray, jnp.ndarray] (Tuple (X, Y) where X is (B, L, D) and Y is (B, P, D)).

masked_mae(y_true, y_pred, mask=None) -> jnp.ndarray

chronax.models.tsmixer.masked_mae

Calculates the Mean Absolute Error, optionally applying a mask.

Parameter	Type	Default	Description
y_true	jnp.ndarray	-	True values (B, P, D).
y_pred	jnp.ndarray	-	Predicted values (B, P, D).
mask	Optional[jnp.ndarray]	None	Optional boolean mask (B, P, D) where False indicates values to ignore.

Returns: jnp.ndarray (Scalar MAE loss).

masked_mse(y_true, y_pred, mask=None) -> jnp.ndarray

chronax.models.tsmixer.masked_mse

Calculates the Mean Squared Error, optionally applying a mask.

Parameter	Type	Default	Description
y_true	jnp.ndarray	-	True values (B, P, D).
y_pred	jnp.ndarray	-	Predicted values (B, P, D).
mask	Optional[jnp.ndarray]	None	Optional boolean mask (B, P, D) where False indicates values to ignore.

Returns: jnp.ndarray (Scalar MSE loss).