chronax.models.vanillatransformer.vanillatransformer

`vanillatransformer_training`

chronax.models.vanillatransformer.vanillatransformer_training

Window construction and training/predict steps for VanillaTransformer.

`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 raw scale. windows: [B, input_size+h] -> scalar.

Parameter	Type	Default	Description
`model`	`VanillaTransformerNet`	-	(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`	`VanillaTransformerNet`	-	(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, indicating training divergence.)

`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`	`VanillaTransformerNet`	-	(undocumented)
`y`	`jnp.ndarray`	-	(undocumented)
`h`	`int`	-	(undocumented)
`input_size`	`int`	-	(undocumented)

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