deepnpts_training
Window construction and training/predict steps for DeepNPTS. Windows are sampled inside the nnx.scan (the per-step PRNG key is the scanned input), so memory stays bounded to one batch even at NF defaults (max_steps=1000, windows_batch_size=1024). The whole loop is one nnx.scan so it stays jax.vmap-traceable for BaseForecaster.conformity_scores. When batch_norm is enabled the carried model's nnx.BatchStat running statistics update in place through the scan.
build_windows(y, input_size, h)
Return [n_windows, input_size+h] rolling windows; step=1.
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
jnp.ndarray |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
h |
int |
- | (undocumented) |
Returns: jnp.ndarray
Raises: ValueError
forward_loss(model, windows, *, h, input_size, loss_fn=mae)
Forward + point loss in raw scale. windows: [B, input_size+h] -> scalar.
| Parameter | Type | Default | Description |
|---|---|---|---|
model |
DeepNPTSNet |
- | (undocumented) |
windows |
jnp.ndarray |
- | (undocumented) |
h |
int |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
loss_fn |
LossFn |
mae |
(undocumented) |
Returns: jnp.ndarray
train(model, y, *, h, input_size, max_steps, windows_batch_size, lr, seed, loss_fn=mae)
Train model in place via a single nnx.scan. Returns per-step losses.
Window sampling replicates neuralforecast's regime-dependent scheme (_base_model.py training_step): when n_windows < windows_batch_size NF draws windows_batch_size indices WITH replacement; otherwise it takes a without-replacement permutation. The branch is chosen on concrete shapes (Python-static), so it bakes into the scan once. Sampling happens per step inside the scan to keep memory bounded to one batch.
| Parameter | Type | Default | Description |
|---|---|---|---|
model |
DeepNPTSNet |
- | (undocumented) |
y |
jnp.ndarray |
- | (undocumented) |
h |
int |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
max_steps |
int |
- | (undocumented) |
windows_batch_size |
int |
- | (undocumented) |
lr |
optax.ScalarOrSchedule |
- | (undocumented) |
seed |
int |
- | (undocumented) |
loss_fn |
LossFn |
mae |
(undocumented) |
Returns: jnp.ndarray (Per-step losses.)
Raises: RuntimeError
predict_step(model, y, *, h, input_size)
Forecast next h steps from the final input_size of y. Returns (h,).
| Parameter | Type | Default | Description |
|---|---|---|---|
model |
DeepNPTSNet |
- | (undocumented) |
y |
jnp.ndarray |
- | (undocumented) |
h |
int |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
Returns: jnp.ndarray (The forecast of shape (h,).)