chronax.models.dilated_rnn.dilated_rnn

`dilated_rnn_training`

chronax.models.dilated_rnn.dilated_rnn_training

Window construction, LR schedule and JIT-compiled training/predict steps for DilatedRNN. Training runs in SCALED space: each rolling window's insample slice sets the robust (median/MAD) statistics, the target is scaled with those same statistics, and the loss is taken there — which is what neuralforecast does via TemporalNorm (scaler_type="robust") before DilatedRNN.forward. Predictions are inverted with the statistics of the prediction context.

`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

`make_lr_schedule(learning_rate, max_steps, num_lr_decays)`

Reproduce NF's StepLR(step_size=max_steps // num_lr_decays, gamma=0.5).

neuralforecast's DilatedRNN defaults to num_lr_decays=3, i.e. the learning rate halves three times over training — a real part of the recipe, not a detail: without it the model keeps taking full-size steps at the end of the run. Torch's StepLR is a staircase, lr * gamma^floor(step / step_size), which is exactly optax.exponential_decay(..., staircase=True).

num_lr_decays <= 0 disables decay (NF's convention). A callable learning_rate is assumed to be a user-supplied optax schedule and is passed through untouched.

Parameter	Type	Default	Description
`learning_rate`	-	-	(undocumented)
`max_steps`	`int`	-	(undocumented)
`num_lr_decays`	`int`	-	(undocumented)

Returns: The learning rate schedule.

`scaled_forward_loss(model, windows, *, h, input_size, scaler, loss_fn=mae)`

Forward + point loss in SCALED space. windows: [B, input_size+h] -> scalar.

The insample slice sets (shift, scale); the target is scaled with the SAME statistics, so the loss is computed exactly where NF computes it.

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

Returns: jnp.ndarray

`train(model, y, *, h, input_size, max_steps, windows_batch_size, lr, seed, scaler=None, 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 the small benchmark series hit, e.g. ~25 windows for AirlinePassengers at input_size=72, h=24, far below the default batch of 128); 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.

lr may be a scalar or an optax schedule — pass the output of :func:make_lr_schedule to get NF's StepLR behaviour.

Note: batches materializes a [max_steps, windows_batch_size, input_size+h] tensor up front. At NF's DilatedRNN defaults (windows_batch_size=128, max_steps=1000, input_size=72, h=24) that is ~49 MB of float32; raising windows_batch_size much further wants a per-step sampler instead.

Parameter	Type	Default	Description
`model`	`DilatedRNNNet`	-	(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 \\| None`	`None`	(undocumented)
`loss_fn`	`LossFn`	`mae`	(undocumented)

Returns: jnp.ndarray Raises: RuntimeError

`predict_step(model, y, *, h, input_size, scaler)`

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

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

Returns: jnp.ndarray