DilatedRNN

chronax.models.DilatedRNN · inherits BaseForecaster

Univariate DilatedRNN forecaster (JAX/Flax-NNX port of neuralforecast.DilatedRNN).

DilatedRNN (Chang et al. 2017, "Dilated Recurrent Neural Networks") stacks recurrent layers that each read the sequence subsampled at their own DILATION rate. A layer at rate r splits the window into r interleaved subsequences and runs them as extra batch rows, so its receptive field spans r timesteps per recurrent step at unchanged cost; stacking rates 1, 2, 4, 8 gives an exponentially growing receptive field with a linear number of steps, which is the point of the architecture (and its answer to vanishing gradients over long windows). Layers are organised into GROUPS (dilations, default [[1, 2], [4, 8]]), with a residual connection added between groups. The final hidden sequence is mapped from the lookback length to the horizon by a single linear "context adapter" (Linear(input_size -> h)), then decoded pointwise by an MLP — i.e. the model is direct multi-step, not recursive. Each rolling window is scaled by a robust median/MAD scaler (NF's scaler_type="robust" default), the model trains in that scaled space with Optax adam under NF's halving learning-rate staircase (num_lr_decays), and predictions are inverted with the prediction context's statistics. float32 throughout.

Maintenance status: Active univariate forecaster. Integrates with the BaseForecaster interface, including conformal prediction intervals via predict(level=...), pickle round-trip, and forecast(fitted=True).

__init__(self, h, input_size=-1, cell_type='LSTM', dilations=((1, 2), (4, 8)), encoder_hidden_size=128, context_size=10, decoder_hidden_size=128, decoder_layers=2, max_steps=1000, learning_rate=0.001, num_lr_decays=3, windows_batch_size=128, scaler_type='robust', random_seed=1, alias='DilatedRNN', loss='mae')

Initialize a DilatedRNN forecaster.

Stores hyperparameters; the network is built lazily at fit time so construction is cheap and side-effect free. Defaults match neuralforecast.DilatedRNN (cell_type="LSTM", dilations=[[1,2],[4,8]], encoder_hidden_size=128, decoder_hidden_size=128, decoder_layers=2, max_steps=1000, learning_rate=1e-3, num_lr_decays=3, windows_batch_size=128, scaler_type="robust"). input_size=-1 resolves to 3 * h.

Parameter	Type	Default	Description
h	int	-	Forecast horizon.
input_size	int	-1	Lookback length; -1 (default) uses 3 * h.
cell_type	str	"LSTM"	One of "GRU", "RNN", "LSTM", "ResLSTM", "AttentiveLSTM".
dilations	Sequence[Sequence[int]]	((1, 2), (4, 8))	Groups of dilation rates. Each inner sequence is one group of stacked layers; a residual connection is added between groups.
encoder_hidden_size	int	128	Recurrent hidden width (all layers).
context_size	int	10	Accepted for signature parity with neuralforecast, which stores it on the model but never reads it in DilatedRNN.forward. It has no effect here either; kept so configs transfer unchanged.
decoder_hidden_size	int	128	Hidden width of the MLP decoder.
decoder_layers	int	2	Total layers in the MLP decoder (1 = bare linear).
max_steps	int	1000	Number of Adam steps.
learning_rate	Union[float, Callable[[int], float]]	1e-3	Scalar, or an optax.ScalarOrSchedule callable. A callable is used as-is and num_lr_decays is ignored.
num_lr_decays	int	3	Number of times the learning rate halves, evenly spread over max_steps (NF's StepLR). <= 0 disables decay.
windows_batch_size	int	128	Rolling windows sampled per step.
scaler_type	str	"robust"	Per-window scaler — "robust" (default), "standard" or "identity".
random_seed	int	1	Seed for parameter init and window sampling.
alias	str	"DilatedRNN"	Display name for external reporting.
loss	Union[str, LossFn]	"mae"	Registry name ("mae"/"mse"/"huber") or a callable.

Raises: ValueError: On an unknown cell_type/scaler_type, an empty or non-positive dilations spec, or decoder_layers < 1.

fit(self, y, X=None) -> Self

Fit the network on a 1-D series.

Builds the network and runs max_steps Adam steps over rolling windows of length input_size + h, sampled per step the way neuralforecast does, under the halving learning-rate staircase.

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D series of length >= input_size + h.
X	jnp.ndarray | None	None	Reserved for exogenous regressors; must be None.

Returns: Self (the fitted forecaster; sets self.model_). Raises: NotImplementedError: If X is provided. ValueError: If y is not 1-D or is shorter than input_size + h. RuntimeError: If a non-finite training loss is observed (divergence).

predict(self, h, X=None, level=None) -> dict

Forecast h steps from the fitted context.

Parameter	Type	Default	Description
h	int	-	Forecast horizon; must satisfy 1 <= h <= self.h (the model is direct-decoded for self.h steps and sliced).
X	jnp.ndarray | None	None	Reserved for exogenous regressors; ignored.
level	list[int | float] | None	None	Optional confidence levels (e.g. [80, 95]). When set, returns conformal lo-XX/hi-XX keys via the inherited BaseForecaster path and requires self.conformal_params. Each call re-fits the model per CV window under vmap — expect minutes.

Returns: dict ({"mean": jnp.ndarray of shape (h,)} plus interval keys when level is provided.) Raises: RuntimeError: If called before fit. ValueError: If h < 1 or h > self.h, or if level is given without self.conformal_params set.

forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> dict

Stateless fit-then-predict on y.

Equivalent to self.fit(y).predict(h=h, level=level), optionally adding a "fitted" key with one-step-ahead in-sample predictions.

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D training series.
h	int	-	Forecast horizon (<= self.h).
X	jnp.ndarray | None	None	Reserved for exogenous regressors; must be None.
X_future	jnp.ndarray | None	None	Reserved for exogenous regressors; must be None.
level	list[int | float] | None	None	Optional confidence levels; see predict.
fitted	bool	False	If True, include "fitted" — one-step-ahead values over the training series, NaN for the first input_size entries.

Returns: dict ({"mean": ..., optional "fitted": ...}). Raises: NotImplementedError: If X or X_future is provided. ValueError / RuntimeError: Forwarded from fit / predict.