TSMixerConfig

chronax.model.TSMixerConfig

Hyperparameters for :class:TSMixer.

__init__(self, h=12, input_size=36, n_series=1, n_block=2, ff_dim=64, dropout=0.1, revin=True, revin_affine=True, temporal_norm_momentum=0.05, feature_norm_momentum=0.05, use_batchnorm=True, use_global_skip=False)

Initializes TSMixerConfig.

Parameter	Type	Default	Description
h	int	12	forecast horizon.
input_size	int	36	history window length (autoregressive lags).
n_series	int	1	number of time series (channels).
n_block	int	2	number of stacked MixingLayers.
ff_dim	int	64	hidden size of the feature-axis feed-forward network.
dropout	float	0.1	dropout rate applied after each mixing sub-block.
revin	bool	True	if True, apply per-channel RevIN before mixing and invert after the output projection.
revin_affine	bool	True	if True (and revin=True), learn per-channel affine parameters γ and β inside RevIN.
temporal_norm_momentum	float	0.05	BatchNorm momentum for temporal mixing.
feature_norm_momentum	float	0.05	BatchNorm momentum for feature mixing.
use_batchnorm	bool	True	(undocumented)
use_global_skip	bool	False	(undocumented)

TemporalMixing

chronax.model.TemporalMixing

Time-axis MLP with normalisation and residual connection.

__init__(self, dropout, use_batchnorm=True)

Initializes TemporalMixing.

Parameter	Type	Default	Description
dropout	float	-	(undocumented)
use_batchnorm	bool	True	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

(undocumented)

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

FeatureMixing

chronax.model.FeatureMixing

Feature-axis two-layer MLP with normalisation and residual connection.

__init__(self, ff_dim, dropout, use_batchnorm=True)

Initializes FeatureMixing.

Parameter	Type	Default	Description
ff_dim	int	-	(undocumented)
dropout	float	-	(undocumented)
use_batchnorm	bool	True	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

(undocumented)

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

MixingLayer

chronax.model.MixingLayer

One TSMixer block: temporal mixing followed by feature mixing.

__init__(self, ff_dim, dropout, use_batchnorm=True)

Initializes MixingLayer.

Parameter	Type	Default	Description
ff_dim	int	-	(undocumented)
dropout	float	-	(undocumented)
use_batchnorm	bool	True	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

(undocumented)

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

TSMixer

chronax.model.TSMixer

JAX/Flax TSMixer — multivariate time-series forecasting model.

Accepts a batch of multivariate windows [B, L, N] and returns [B, h, N] forecasts. Carries batch_stats (running BatchNorm statistics) as a mutable variable collection.

__init__(self, config)

Initializes TSMixer.

Parameter	Type	Default	Description
config	TSMixerConfig	-	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

(undocumented)

Parameters:

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).