train.py

Training utilities for Chronax TiDE. Pure JAX/Flax/Optax — no PyTorch, no numpy at training time.

Two training paths are provided:

Window-based (train_step_windows / eval_step_windows): Primary path. Accepts raw [B, input_size + h] windows, applies per-window RobustScaler (median/MAD) internally, then optimises the supplied loss. Used by chronax.models.tide.forecaster.TiDEForecaster and mirrors the NBEATS training path exactly.
Batch-based (train_step / eval_step): Compatibility path for pre-assembled {insample_y, outsample_y, ...} dicts — required when exogenous variables are present.

`TrainState`

train.TrainState · inherits flax.training.train_state.TrainState

Standard Flax TrainState; aliased for forward compatibility.

`create_train_state`

train.create_train_state(rng, config, *, learning_rate=1e-3, grad_clip=1.0, cosine_decay_steps=0, warmup_steps=0, num_lr_decays=-1, max_steps=1000, weight_decay=0.0, optimizer=None)

Initialise model parameters and optimiser state.

Parameter	Type	Default	Description
`rng`	`jax.Array`	-	PRNG key for parameter initialisation.
`config`	`TiDEConfig`	-	Model architecture config.
`learning_rate`	`float`	`1e-3`	Adam(W) base learning rate.
`grad_clip`	`float`	`1.0`	Global gradient-norm clip (0 = disabled).
`cosine_decay_steps`	`int`	`0`	Warmup-cosine schedule length (0 = disabled).
`warmup_steps`	`int`	`0`	Linear warmup steps for the cosine schedule.
`num_lr_decays`	`int`	`-1`	NF-style StepLR decays (-1 = constant LR).
`max_steps`	`int`	`1000`	Total training steps (used for StepLR spacing).
`weight_decay`	`float`	`0.0`	Use AdamW when > 0; plain Adam otherwise.
`optimizer`	`Optional[optax.GradientTransformation]`	`None`	Pre-built transform that overrides all defaults.

Returns: TrainState

`train_step`

train.train_step(state, batch, rng, loss_fn=masked_mae)

JIT-compiled training step.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`batch`	`Dict[str, Optional[jnp.ndarray]]`	-	(undocumented)
`rng`	`jax.Array`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

Returns: Tuple[TrainState, jnp.ndarray, jnp.ndarray] (Returns (new_state, loss, predictions).)

`eval_step`

train.eval_step(state, batch, loss_fn=masked_mae)

JIT-compiled evaluation step. Deterministic, no gradients.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`batch`	`Dict[str, Optional[jnp.ndarray]]`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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

`train_loop`

train.train_loop(state, train_batches, *, num_epochs=10, rng=None, eval_batches=None, loss_fn=masked_mae)

Drive train_step / eval_step over multiple epochs.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`train_batches`	`Iterable[Dict[str, Optional[jnp.ndarray]]]`	-	(undocumented)
`num_epochs`	`int`	`10`	(undocumented)
`rng`	`Optional[jax.Array]`	`None`	(undocumented)
`eval_batches`	`Optional[Iterable[Dict[str, Optional[jnp.ndarray]]]]`	`None`	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

Returns: Tuple[TrainState, List[Dict[str, float]]]

`train_step_windows`

train.train_step_windows(state, windows, rng, input_size, loss_fn=masked_mae)

JIT-compiled window-based training step with per-window RobustScaler.

Parameter	Type	Default	Description
`state`	`TrainState`	-	Current `TrainState`.
`windows`	`jnp.ndarray`	-	Raw `[B, input_size + h]` windows (unscaled).
`rng`	`jax.Array`	-	PRNG key for dropout.
`input_size`	`int`	-	History length `L` (static; triggers re-trace when changed).
`loss_fn`	`Callable`	`masked_mae`	`(y_true, y_pred) -> scalar` (default: `masked_mae`).

Returns: Tuple[TrainState, jnp.ndarray, jnp.ndarray] (Returns (new_state, loss, predictions).)

`eval_step_windows`

train.eval_step_windows(state, windows, input_size, loss_fn=masked_mae)

JIT-compiled window-based evaluation step. Deterministic, no gradients.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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

`train_step_windows_std`

train.train_step_windows_std(state, windows, rng, input_size, loss_fn=masked_mae)

JIT-compiled window-based training step with per-window mean/std scaler.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`rng`	`jax.Array`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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

`eval_step_windows_std`

train.eval_step_windows_std(state, windows, input_size, loss_fn=masked_mae)

JIT-compiled window-based eval step with per-window mean/std scaler.

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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

`train_step_windows_raw`

train.train_step_windows_raw(state, windows, rng, input_size, loss_fn=masked_mae)

JIT-compiled training step for pre-normalised windows (no per-window scaler).

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`rng`	`jax.Array`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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

`eval_step_windows_raw`

train.eval_step_windows_raw(state, windows, input_size, loss_fn=masked_mae)

JIT-compiled eval step for pre-normalised windows (no per-window scaler).

Parameter	Type	Default	Description
`state`	`TrainState`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`loss_fn`	`Callable`	`masked_mae`	(undocumented)

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