TrainState

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

Extends Flax TrainState with a batch_stats field for BatchNorm.

Keeping batch_stats separate from params ensures the Adam optimizer only tracks and updates the trainable parameters; running statistics are updated via exponential moving average inside train_step, never via gradient descent.

create_train_state

chronax.train.create_train_state

Initialise model parameters, BatchNorm statistics, and optimiser state.

Parameter	Type	Default	Description
rng	jax.Array	-	PRNG key for parameter initialisation.
config	TSMixerConfig	-	model architecture config.
learning_rate	float	1e-3	peak learning rate.
weight_decay	float	0.0	if > 0, use AdamW instead of Adam.
grad_clip	float	1.0	global gradient-norm clip threshold (0 = disabled).
cosine_decay_steps	int	0	if > 0, use warmup + cosine decay schedule.
warmup_steps	int	0	linear warmup steps prepended to the cosine decay.
optimizer	Optional[optax.GradientTransformation]	None	optional pre-built optax transform; overrides all defaults.

Returns: TrainState

train_step

chronax.train.train_step

JIT-compiled training step.

Differentiates only w.r.t. state.params (trainable weights); state.batch_stats is updated via the running-average returned by model.apply(..., mutable=['batch_stats']).

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

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

eval_step

chronax.train.eval_step

JIT-compiled evaluation step. Uses running BatchNorm statistics.

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

Returns: Tuple[jnp.ndarray, jnp.ndarray] (loss, output).

scan_train_loop

chronax.train.scan_train_loop

Train for S steps via jax.lax.scan.

Note: does not update batch_stats (BatchNorm running statistics). Use the Python train_step loop when the model uses BatchNorm.

Parameter	Type	Default	Description
state	TrainState	-	(undocumented)
all_ins	jnp.ndarray	-	(undocumented)
all_out	jnp.ndarray	-	(undocumented)
all_rngs	jnp.ndarray	-	(undocumented)

Returns: Tuple[TrainState, jnp.ndarray]

train_loop

chronax.train.train_loop

Drive train_step / eval_step over multiple epochs.

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

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