itransformer_training
Window construction and JIT-compiled training/predict steps for iTransformer.
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
forward_loss(model, windows, *, h, input_size, loss_fn=mae)
Forward + point loss in ORIGINAL scale (RevIN denorms inside the net).
| Parameter | Type | Default | Description |
|---|---|---|---|
model |
ITransformerNet |
- | (undocumented) |
windows |
jnp.ndarray |
- | [B, input_size+h] -> scalar. The univariate series carries a channel dim of 1, so the insample window is reshaped to [B, L, 1]. |
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.
The whole loop is one nnx.scan (carry = (model, optimizer)), which keeps the function jax.vmap-traceable for BaseForecaster.conformity_scores. 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 (oversampling with duplicates — the regime small benchmark series hit, e.g. ~25 windows for AirlinePassengers at input_size=72, h=24, both << the batch size); otherwise it takes a without-replacement permutation of windows_batch_size windows. Getting this branch right is load-bearing for accuracy parity, so we do NOT collapse it to a full batch.
Note: batches materializes a [max_steps, windows_batch_size, input_size+h] tensor up front. At iTransformer's defaults (windows_batch_size=32) this is small; if you raise windows_batch_size substantially, reduce max_steps or sample per-step instead to bound memory.
| Parameter | Type | Default | Description |
|---|---|---|---|
model |
ITransformerNet |
- | (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 (If non-finite loss is detected during training.)
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 |
ITransformerNet |
- | (undocumented) |
y |
jnp.ndarray |
- | (undocumented) |
h |
int |
- | (undocumented) |
input_size |
int |
- | (undocumented) |
Returns: jnp.ndarray (Shape (h,).)