chronax.models.gru.gru

`build_windows`

gru_training.build_windows

Return [n_windows, input_size+h] rolling windows; step=1.

`build_windows(y, input_size, h)`

Parameter	Type	Default	Description
`y`	`jnp.ndarray`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`h`	`int`	-	(undocumented)

Returns: jnp.ndarray

`scaled_forward_loss`

gru_training.scaled_forward_loss

Forward + point-loss in scaled space. windows: [B, input_size+h] -> scalar.

loss_fn must be a callable with signature (pred, target) -> scalar; see :mod:chronax.models.gru.gru_losses for the built-in registry.

`scaled_forward_loss(model, windows, *, h, input_size, scaler, loss_fn=mae)`

Parameter	Type	Default	Description
`model`	`GRUNet`	-	(undocumented)
`windows`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`scaler`	`Scaler`	-	(undocumented)
`loss_fn`	`LossFn`	`mae`	(undocumented)

Returns: jnp.ndarray

`train`

gru_training.train

Train model in place. Returns per-step training losses.

The training loop runs as a single nnx.scan, which makes the entire function jax.vmap-traceable. Mutable Flax NNX state (model parameters, optimizer momentum) flows through the scan as nnx.Carry — NNX handles the graph-vs-state separation internally.

Per-step batches are pre-sampled outside the scan so the body is a pure function of (model, optimizer, batch) -> mutated state, loss. The non-finite-loss check is one host-side reduction over the loss array after the scan returns, rather than a per-step sync.

`train(model, y, *, h, input_size, max_steps, batch_size, lr, seed, scaler=None, loss_fn=mae)`

Parameter	Type	Default	Description
`model`	`GRUNet`	-	(undocumented)
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`max_steps`	`int`	-	(undocumented)
`batch_size`	`int`	-	(undocumented)
`lr`	`optax.ScalarOrSchedule`	-	(undocumented)
`seed`	`int`	-	(undocumented)
`scaler`	`Scaler \\| None`	`None`	(undocumented)
`loss_fn`	`LossFn`	`mae`	(undocumented)

Returns: jnp.ndarray (per-step training losses) Raises: RuntimeError

`predict_step`

gru_training.predict_step

Forecast next h steps from final input_size of y. Returns shape (h,).

Delegates to module-level _jit_forward_deterministic so repeated calls on the same model do not re-trace.

`predict_step(model, y, *, h, input_size, scaler)`

Parameter	Type	Default	Description
`model`	`GRUNet`	-	(undocumented)
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)
`scaler`	`Scaler`	-	(undocumented)

Returns: jnp.ndarray