chronax.models.timemixer.timemixer

TimeMixer forecaster: BaseForecaster wrapper around the flax.nnx backbone.

Univariate wrapper (n_series=1) over the N-generic multiscale-mixing network. Point losses; intervals via the conformal path. Normalization lives inside the network (per-scale RevIN), so windows train in original scale; optional Box-Cox variance stabilization mirrors the fleet's use_boxcox convention (off = faithful port). float32 throughout.

`TimeMixer`

chronax.models.timemixer_model.TimeMixer · inherits BaseForecaster

TimeMixer: Decomposable Multiscale Mixing (flax.nnx port of neuralforecast.TimeMixer).

Wang et al., 2024 — https://openreview.net/pdf?id=7oLshfEIC2. The input window is repeatedly downsampled into a scale pyramid; each scale is instance-normalized and season/trend-decomposed, seasons mix bottom-up and trends top-down through small Linear+GELU chains (Past-Decomposable Mixing), and per-scale predictors map every scale onto the horizon before the outputs are summed (Future-Multipredictor Mixing). Channel-dependent (default) and channel-independent modes are both supported, as are moving-average and DFT decompositions and avg/max/conv downsampling. Exogenous inputs are not modeled. Point losses; conformal intervals. float32 throughout.

Attributes: * uses_exog: bool (False) * alias: str (TimeMixer) * conformal_params: ConformalIntervals | None * model_: TimeMixerNet | None

`init(self, h, input_size=-1, d_model=32, d_ff=32, dropout=0.1, e_layers=4, top_k=5, decomp_method='moving_avg', moving_avg=25, channel_independence=0, down_sampling_layers=1, down_sampling_window=2, down_sampling_method='avg', use_norm=True, decoder_input_size_multiplier=0.5, use_boxcox=False, loss='mae', max_steps=1000, learning_rate=1e-3, windows_batch_size=32, random_seed=1, alias='TimeMixer')`

(undocumented)

Parameter	Type	Default	Description
`h`	-	-	(undocumented)
`input_size`	`int`	-1	(undocumented)
`d_model`	`int`	32	(undocumented)
`d_ff`	`int`	32	(undocumented)
`dropout`	`float`	0.1	(undocumented)
`e_layers`	`int`	4	(undocumented)
`top_k`	`int`	5	(undocumented)
`decomp_method`	`str`	"moving_avg"	(undocumented)
`moving_avg`	`int`	25	(undocumented)
`channel_independence`	`int`	0	(undocumented)
`down_sampling_layers`	`int`	1	(undocumented)
`down_sampling_window`	`int`	2	(undocumented)
`down_sampling_method`	`str`	"avg"	(undocumented)
`use_norm`	`bool`	True	(undocumented)
`decoder_input_size_multiplier`	`float`	0.5	(undocumented)
`use_boxcox`	`bool`	False	(undocumented)
`loss`	`str`	"mae"	(undocumented)
`max_steps`	`int`	1000	(undocumented)
`learning_rate`	`float`	0.001	(undocumented)
`windows_batch_size`	`int`	32	(undocumented)
`random_seed`	`int`	1	(undocumented)
`alias`	`str`	"TimeMixer"	(undocumented)

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

(undocumented)

Parameters:

Parameter	Type	Default	Description
`y`	-	-	(undocumented)
`X`	-	None	(undocumented)

Returns: Self (the fitted forecaster; sets self.model_). Raises: * NotImplementedError: If X is not None. * ValueError: If y is not 1-D, too short for input_size, or contains non-positive values when use_boxcox=True.

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

Generates a forecast for horizon h based on the fitted model and training data.

Parameters:

Parameter	Type	Default	Description
`h`	-	-	(undocumented)
`X`	-	None	(undocumented)
`level`	-	None	(undocumented)

Returns: dict Keys: {"mean": jnp.ndarray}. If level is provided and self.conformal_params is set, includes lower_bound and upper_bound. Raises: * RuntimeError: If fit(y) has not been called. * ValueError: If h is not positive, h exceeds the trained horizon (self.h), or level is provided but model.conformal_params is missing.

`conformity_scores(self, y, X=None)`

Walk-forward conformity scores; with use_boxcox=True the CV re-fits reuse the lambda frozen by the eager fit (selection cannot run on traced windows).

Parameters:

Parameter	Type	Default	Description
`y`	-	-	(undocumented)
`X`	-	None	(undocumented)

Raises: * ValueError: If use_boxcox=True requires a fitted estimator before conformal CV.

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

Stateless fit-then-predict.

Parameters:

Parameter	Type	Default	Description
`y`	-	-	(undocumented)
`h`	-	-	(undocumented)
`X`	-	None	(undocumented)
`X_future`	-	None	(undocumented)
`level`	-	None	(undocumented)
`fitted`	`bool`	False	(undocumented)

Returns: dict Keys: {"mean": jnp.ndarray}. If level is provided, includes lower_bound and upper_bound. If fitted=True, includes fitted: jnp.ndarray. Raises: * NotImplementedError: If X or X_future is not None.