GRU

chronax.models.GRU · inherits BaseForecaster

Provides a JAX/Flax/Optax GRU forecaster wrapper with a stable class interface for fitting and forecasting univariate time series.

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

Attributes: * uses_exog: False * alias: Display name for external reporting. * model_: Fitted network after fit (GRUNet | None). * conformal_params: Conformal calibration set by BaseForecaster when fit is given prediction_intervals; consumed by predict(level=...).

__init__(self, h: int, input_size: int = -1, hidden_size: int = 200, n_layers: int = 2, decoder_hidden_size: int = 128, dropout: float = 0.0, max_steps: int = 1000, learning_rate: Union[float, Callable[[int], float]] = 1e-3, batch_size: int = 128, random_seed: int = 1, alias: str = 'GRU', loss: Union[str, LossFn] = 'mae', recurrent_init: str = 'uniform')

Stores hyperparameters; the network is not built until fit is called (so __init__ is cheap and side-effect free).

Parameter	Type	Default	Description
h	int	-	Forecast horizon.
input_size	int	-1	Input-window length. Use -1 for 3 * h.
hidden_size	int	200	Encoder hidden dimension.
n_layers	int	2	Number of stacked GRU cells.
decoder_hidden_size	int	128	MLP-decoder hidden dimension.
dropout	float	0.0	Inter-layer dropout rate.
max_steps	int	1000	Optimizer steps for fit.
learning_rate	Union[float, Callable[[int], float]]	1e-3	Adam learning rate. Accepts either a scalar or any optax.ScalarOrSchedule (a callable mapping step → LR, e.g. optax.cosine_decay_schedule). If you pass a callable and later want to pickle the fitted estimator, ensure the callable itself is picklable (a class-based callable or a module-level helper — closures returned by some optax helpers are not).
batch_size	int	128	Windows per training step.
random_seed	int	1	Random seed.
alias	str	"GRU"	User-facing model name.
loss	Union[str, LossFn]	"mae"	Either a registered name ("mae", "mse", "huber") from :mod:chronax.models.gru.gru_losses or a callable with signature (pred, target) -> scalar. Strings are validated lazily at fit time so the constructor stays cheap and the resulting estimator pickles cleanly.
recurrent_init	str	"uniform"	Initializer for the GRU recurrent kernels. "uniform" (default) samples from Uniform(-1/sqrt(H), 1/sqrt(H)); "orthogonal" uses Saxe-style orthogonal initialization, which stabilises gradients through time and often helps with longer input_size. The input kernel and biases use the uniform init regardless.

fit(self, y: jnp.ndarray, X: jnp.ndarray | None = None) -> Self

Fit the GRU on a univariate series.

Builds the network with the configured hyperparameters, then runs self.max_steps Adam steps over rolling windows of length input_size + h sampled uniformly at random. Each window is normalized with the per-window robust scaler before the forward pass; the configured loss is applied in scaled space.

Parameters:

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

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

predict(self, h: int, X: jnp.ndarray | None = None, level: list[int | float] | None = None) -> dict

Forecast from the fitted context.

Runs a single deterministic forward pass on the cached context window (last input_size of the fit-time series), then slices to the requested h. Always materializes self.h outputs internally so the JIT cache is shared across calls regardless of the caller's h.

Parameters:

Parameter	Type	Default	Description
h	int	-	Forecast horizon. Must satisfy h <= self.h.
X	jnp.ndarray | None	None	Reserved for future exogenous regressors; ignored in v1.
level	list[int | float] | None	None	If provided, returns conformal prediction intervals as additional lo-XX/hi-XX keys (e.g. lo-80, hi-80 for level=[80]). Requires self.conformal_params to be set. The inherited conformity-score computation re-fits the model per CV window; on GRU expect MINUTES per call. Reduce n_windows or max_steps if interactive feedback matters.

Returns: dict ({"mean": jnp.ndarray of shape (h,)}. If level is provided, includes lo-XX and hi-XX keys.) Raises: * ValueError: If h > self.h. * RuntimeError: If called before fit. * ValueError: If level is provided but self.conformal_params is not set. * RuntimeError: If level is provided but self._train_y is None.

forecast(self, y: jnp.ndarray, h: int, X: jnp.ndarray | None = None, X_future: jnp.ndarray | None = None, level: list[int | float] | None = None, fitted: bool = False) -> dict

Stateless fit-then-predict.

Equivalent to self.fit(y).predict(h=h, level=level) (plus an extra "fitted" key when fitted=True) and numerically identical to that pattern when seeded the same way.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D training series.
h	int	-	Forecast horizon (<= self.h).
X	jnp.ndarray | None	None	Reserved; must be None in v1.
X_future	jnp.ndarray | None	None	Reserved; must be None in v1.
level	list[int | float] | None	None	Optional confidence levels (e.g. [80, 95]). When set, returns conformal intervals via the inherited BaseForecaster.add_confidence_intervals path. See note on cost (minutes per call) in predict.
fitted	bool	False	If True, the returned dict additionally contains a "fitted" key with one-step-ahead predictions over the training series. The first input_size entries are NaN (no input window available); the remaining entries are finite. Matches the convention used by other Chronax models.

Returns: dict ({"mean": ..., optional "fitted": ...}.) Raises: * NotImplementedError: If X or X_future is non-None. * ValueError: Forwarded from fit / predict. * RuntimeError: Forwarded from fit (training divergence) or predict.