softs_training.py

Window construction and JIT-compiled training/predict steps for SOFTS.

build_windows

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

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)

forward_loss

forward_loss(model: SOFTSNet, windows: jnp.ndarray, *, h: int, input_size: int, loss_fn: LossFn = mae) -> jnp.ndarray

Forward + point loss in ORIGINAL scale (RevIN denorms inside the net).

windows: [B, input_size+h] -> scalar. The univariate series carries a channel dim of 1, so the insample window is reshaped to [B, L, 1].

Parameter	Type	Default	Description
model	SOFTSNet	-	(undocumented)
windows	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
input_size	int	-	(undocumented)
loss_fn	LossFn	mae	(undocumented)

train

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

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.

STAD's stochastic pooling and the dropout layers both draw from the model's nnx.Rngs; because the model is the scan carry, those key streams advance per step without any explicit threading here.

Note: batches materializes a [max_steps, windows_batch_size, input_size+h] tensor up front. At SOFTS'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	SOFTSNet	-	(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 Raises: RuntimeError: If non-finite loss is detected during training.

predict_step

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

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

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

Returns: jnp.ndarray