chronax.models.patchtst.patchtst

`PatchTST`

chronax.models.patchtst_model.PatchTST · inherits BaseForecaster

Univariate PatchTST forecaster (JAX/Flax-NNX port of neuralforecast.PatchTST). The encoder patches the input window, embeds each patch, and runs a stack of transformer layers (residual attention + BatchNorm + GELU feed-forward) with RevIN normalization applied inside the network. Trained in original scale with Optax adam and a pluggable point loss. 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).

Attributes

Attribute	Type	Description
`uses_exog`	`bool`	`False`
`alias`	`str`	`"PatchTST"` (default)
`conformal_params`	`Any`	Parameters used for conformal prediction.
`model_`	`PatchTSTNet` or `None`	The fitted Flax-NNX network module.

`init(self, h, input_size=-1, patch_len=16, stride=8, hidden_size=128, n_heads=16, encoder_layers=3, linear_hidden_size=256, dropout=0.2, fc_dropout=0.2, head_dropout=0.0, attn_dropout=0.0, activation='gelu', revin=True, revin_affine=False, revin_subtract_last=True, max_steps=5000, learning_rate=1e-4, windows_batch_size=1024, random_seed=1, alias='PatchTST', loss='mae')`

Initialize a PatchTST forecaster. Stores hyperparameters; the network is built lazily at fit time so construction is cheap and side-effect free. Defaults match neuralforecast.PatchTST. input_size=-1 resolves to 3 * h. loss is a registry name ("mae"/"mse"/"huber") or a callable; learning_rate is a scalar or an optax.ScalarOrSchedule; activation is "gelu" or "relu". If the fitted estimator will be pickled, any callable passed for loss/learning_rate must itself be picklable (a class-based callable or module-level function).

Parameter	Type	Default	Description
`h`	`int`	-	(undocumented)
`input_size`	`int`	`-1`	Resolves to `3 * h` if `< 1`.
`patch_len`	`int`	`16`	(undocumented)
`stride`	`int`	`8`	(undocumented)
`hidden_size`	`int`	`128`	(undocumented)
`n_heads`	`int`	`16`	(undocumented)
`encoder_layers`	`int`	`3`	(undocumented)
`linear_hidden_size`	`int`	`256`	(undocumented)
`dropout`	`float`	`0.2`	(undocumented)
`fc_dropout`	`float`	`0.2`	(undocumented)
`head_dropout`	`float`	`0.0`	(undocumented)
`attn_dropout`	`float`	`0.0`	(undocumented)
`activation`	`str`	`"gelu"`	`"gelu"` or `"relu"`.
`revin`	`bool`	`True`	(undocumented)
`revin_affine`	`bool`	`False`	(undocumented)
`revin_subtract_last`	`bool`	`True`	(undocumented)
`max_steps`	`int`	`5000`	(undocumented)
`learning_rate`	`Union[float, Callable[[int], float]]`	`1e-4`	Scalar or an `optax.ScalarOrSchedule`. Must be picklable if callable.
`windows_batch_size`	`int`	`1024`	(undocumented)
`random_seed`	`int`	`1`	(undocumented)
`alias`	`str`	`"PatchTST"`	(undocumented)
`loss`	`Union[str, LossFn]`	`"mae"`	Registry name (`"mae"`/`"mse"`/`"huber"`) or a callable. Must be picklable if callable.

`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.

Parameters:

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.

Parameters:

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` — on a full PatchTST this costs 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.

Parameters:

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.