build_windows

chronax.kan_training.build_windows

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.

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

Returns: tuple[jnp.ndarray, jnp.ndarray] (windows, mask). Raises:

Exception	Description
ValueError	If series length is too short for input_size and h.

scaled_forward_loss

chronax.kan_training.scaled_forward_loss

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.

Parameter	Type	Default	Description
model	KANNet	-	(undocumented)
windows	jnp.ndarray	-	(undocumented)
mask	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
input_size	int	-	(undocumented)
scaler	Scaler	-	(undocumented)
loss_fn	LossFn	mae	(undocumented)

Returns: jnp.ndarray

train

chronax.kan_training.train

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).

Parameter	Type	Default	Description
model	KANNet	-	(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)
scaler	Scaler	-	(undocumented)
loss_fn	LossFn	mae	(undocumented)

Returns: jnp.ndarray (per-step losses). Raises:

Exception	Description
RuntimeError	If non-finite loss is detected, indicating training divergence. Consider lowering learning_rate, reducing windows_batch_size, or using scaler='robust'.

predict_step

chronax.kan_training.predict_step

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

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

Returns: jnp.ndarray (shape (h,)).