TrainState
chronax.train.TrainState · inherits flax.training.train_state.TrainState
Standard Flax TrainState; aliased for forward compatibility.
create_train_state
chronax.train.create_train_state
Initialise model parameters and optimiser state.
| Parameter | Type | Default | Description |
|---|---|---|---|
rng |
jax.Array |
- | PRNG key used for parameter initialisation only. |
config |
RNNConfig |
- | model architecture config. |
learning_rate |
float |
1e-3 |
peak learning rate for the default Adam(W) optimiser. |
weight_decay |
float |
0.0 |
if > 0 use optax.adamw, else optax.adam. |
grad_clip |
float |
1.0 |
global gradient norm clipping threshold (0 = disabled). |
cosine_decay_steps |
int |
0 |
if > 0 use a warmup + cosine decay schedule that decays from learning_rate to learning_rate * 0.01 over this many steps. |
warmup_steps |
int |
0 |
linear warmup steps at the start of 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 using jax.value_and_grad.
| 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] (new_state, loss, predictions).
eval_step
chronax.train.eval_step
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_step_recurrent
chronax.train.train_step_recurrent
JIT-compiled training step for recurrent mode. Uses a 1-step-ahead prediction loss, matching NeuralForecast RNN's h_train=1 training objective. The model returns output[:, -h:] where each position i predicts y[i+1] (one step ahead). next_step_y [B, h, 1] is built as concat(insample_y[:, -(h-1):], outsample_y[:, :1]) so every comparison is a valid next-step target.
| Parameter | Type | Default | Description |
|---|---|---|---|
state |
TrainState |
- | (undocumented) |
insample_y |
jnp.ndarray |
- | (undocumented) |
next_step_y |
jnp.ndarray |
- | (undocumented) |
rng |
jax.Array |
- | (undocumented) |
loss_fn |
Callable |
masked_mae |
(undocumented) |
Returns: Tuple[TrainState, jnp.ndarray, jnp.ndarray]
eval_step_recurrent
chronax.train.eval_step_recurrent
JIT-compiled eval step for recurrent mode. Deterministic, no gradients.
| Parameter | Type | Default | Description |
|---|---|---|---|
state |
TrainState |
- | (undocumented) |
insample_y |
jnp.ndarray |
- | (undocumented) |
next_step_y |
jnp.ndarray |
- | (undocumented) |
loss_fn |
Callable |
masked_mae |
(undocumented) |
Returns: Tuple[jnp.ndarray, jnp.ndarray]
train_loop
chronax.train.train_loop
Drive train_step / eval_step over multiple epochs. train_batches may be a list, generator factory, or any iterable. If it is a generator object that gets exhausted after one pass, you should materialise it (e.g. into a list) before calling this function.
| 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]]]