nlinear_training.py Window construction and JIT-compiled training/predict steps for NLinear.

build_windows

build_windows(y: jnp.ndarray, input_size: int, h: int) -> tuple[jnp.ndarray, jnp.ndarray]

Rolling training windows with neuralforecast-style right-padding.

Right-pads y with h zeros before unfolding (NF padder_train = ConstantPad1d((0, h))), yielding len(y) - input_size windows of length input_size + h — including ~h partial-horizon windows whose context reaches the end of the series. Returns (windows, mask) where mask is 1 on real points and 0 on the padded tail, so the loss can drop padded horizon steps. The insample (first input_size) of every window is fully real. Returns shape [n_windows, input_size+h] each.

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

scaled_forward_loss

scaled_forward_loss(model: NLinearNet, windows: jnp.ndarray, mask: jnp.ndarray, *, h: int, input_size: int, scaler: Scaler, loss_fn: LossFn = mae) -> jnp.ndarray

Forward + masked point loss in SCALED space. windows/mask: [B, input_size+h] -> scalar.

The insample is always real (padding is target-side only), so the scaler sees real values. Padded horizon steps are excluded via the outsample mask.

Returns: jnp.ndarray

train

train(model: NLinearNet, y: jnp.ndarray, *, h: int, input_size: int, max_steps: int, windows_batch_size: int, lr: optax.ScalarOrSchedule, seed: int, scaler: Scaler, loss_fn: LossFn = mae) -> jnp.ndarray

Train model in place via a single nnx.scan. Returns per-step losses.

Single nnx.scan (carry = (model, optimizer)) keeps train() vmap-traceable for BaseForecaster.conformity_scores. Window sampling replicates neuralforecast's regime-dependent scheme (with-replacement when n_windows < windows_batch_size, else a permutation). The scan iterates an int32 index tensor and gathers windows in-step (avoids materializing a large float32 batch tensor).

Returns: jnp.ndarray

predict_step

predict_step(model: NLinearNet, y: jnp.ndarray, *, h: int, input_size: int, scaler: Scaler) -> jnp.ndarray

Forecast next h steps from the final input_size of y, inverse-scaled. Returns (h,).

Returns: jnp.ndarray