chronax.models.softssharp.softssharp

`softssharp_training`

chronax.models.softssharp.softssharp_training

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

`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: If series length is too short for the specified window size.

`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`	`SOFTSSharpNet`	-	(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. STADSharp's stochastic pooling, its position-encoding gate, and the dropout layers all 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 SOFTSSharp'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`	`SOFTSSharpNet`	-	(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: Non-finite loss detected, indicating training diverged.

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

Returns: jnp.ndarray