train

Training utilities for Chronax N-BEATS.

Pure JAX/Flax/Optax — no PyTorch, no numpy.

TrainState

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

Standard Flax TrainState; aliased for forward compatibility.

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	NBEATSConfig	-	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	Use warmup-cosine schedule decaying over this many steps (0 = disabled; takes precedence over num_lr_decays).
warmup_steps	int	0	Linear warmup steps for the cosine schedule.
num_lr_decays	int	-1	NF-style StepLR decays with gamma=0.5 (used only when cosine_decay_steps <= 0; -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(state, windows, rng, input_size, loss_fn=masked_mae, scale=True)

JIT-compiled window-based training step.

Parameter	Type	Default	Description
state	TrainState	-	Current TrainState.
windows	jnp.ndarray	-	Raw [B, input_size + h] windows.
rng	jax.Array	-	PRNG key for dropout.
input_size	int	-	History length (static; re-compiles when changed).
loss_fn	Callable	masked_mae	(y_true, y_pred) -> scalar (default: masked_mae).
scale	bool	True	Apply per-window RobustScaler when True (static flag).

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

eval_step(state, windows, input_size, loss_fn=masked_mae, scale=True)

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

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

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

train_batch_step(state, batch, rng, loss_fn=masked_mae)

JIT-compiled batch-based training step on pre-scaled inputs.

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]

eval_batch_step(state, batch, loss_fn=masked_mae)

JIT-compiled batch evaluation step. Deterministic, no grads.

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(state, train_batches, *, num_epochs=10, rng=None, eval_batches=None, loss_fn=masked_mae)

Drive train_batch_step / eval_batch_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]]]