stemgnn_training
stemgnn_training
Window construction and scan-based training for StemGNN. Windows are scaled per-window (scaler stats on the insample target), the loss is computed in scaled space, and the whole training loop is one nnx.scan so it stays vmap-traceable for BaseForecaster.conformity_scores. Two RNG streams split off seed: batch-index keys (regime-dependent window sampling) and per-step dropout keys for the attention dropout — shuffling batches never perturbs which attention entries drop on a given step.
build_windows(y, input_size, h)
Rolling windows over y right-padded with h zeros.
The series is padded with h trailing zeros before windowing, so every window with at least one valid target point is kept and the padded tail is masked out of the loss. These partial windows put insample contexts ending at the very last observations into training — on trending series that is where the forecast-relevant regime lives.
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
jnp.ndarray |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
h |
int |
- | (undocumented) |
Returns: tuple[jnp.ndarray, jnp.ndarray] (windows [n, input_size+h], target_mask [n, h]) with n = len(y) - input_size; the mask is 1.0 where the target position is a real observation and 0.0 in the zero-padded tail.
Raises: ValueError
forward_loss(net, y_windows, target_mask=None, *, h, input_size, scaler, loss_fn, dropout_key=None, deterministic=True)
Scale, forward, and reduce a point/quantile loss in scaled space.
target_mask [B, h] marks real target positions (0 in the padded tail); the loss is the masked mean over valid elements. None means all-valid.
| Parameter | Type | Default | Description |
|---|---|---|---|
net |
- | - | (undocumented) |
y_windows |
- | - | (undocumented) |
target_mask |
- | None |
[B, h] marks real target positions (0 in the padded tail); the loss is the masked mean over valid elements. None means all-valid. |
h |
- | - | (undocumented) |
input_size |
- | - | (undocumented) |
scaler |
- | - | (undocumented) |
loss_fn |
- | - | (undocumented) |
dropout_key |
- | None |
(undocumented) |
deterministic |
- | True |
(undocumented) |
train(net, y, *, h, input_size, max_steps, windows_batch_size, lr, num_lr_decays, seed, loss_fn, scaler)
Train net in place via one nnx.scan. Returns per-step losses.
Batch-index sampling splits by regime: with replacement when the dataset has fewer windows than windows_batch_size, without replacement otherwise. Adam runs on the StepLR schedule from _lr_schedule.
| Parameter | Type | Default | Description |
|---|---|---|---|
net |
- | - | (undocumented) |
y |
- | - | (undocumented) |
h |
- | - | (undocumented) |
input_size |
- | - | (undocumented) |
max_steps |
- | - | (undocumented) |
windows_batch_size |
- | - | (undocumented) |
lr |
- | - | (undocumented) |
num_lr_decays |
- | - | (undocumented) |
seed |
- | - | (undocumented) |
loss_fn |
- | - | (undocumented) |
scaler |
- | - | (undocumented) |
Returns: jnp.ndarray (per-step losses).
Raises: RuntimeError
predict_step(net, y_context, *, h, input_size, scaler)
Forecast next h steps from the final input_size of the series.
Returns [h, multiplier] in the original scale.
| Parameter | Type | Default | Description |
|---|---|---|---|
net |
- | - | (undocumented) |
y_context |
- | - | (undocumented) |
h |
- | - | (undocumented) |
input_size |
- | - | (undocumented) |
scaler |
- | - | (undocumented) |
Returns: jnp.ndarray ([h, multiplier]) in the original scale.