VanillaTransformer

chronax.models.vanillatransformer_model.VanillaTransformer · inherits BaseForecaster

Univariate VanillaTransformer forecaster (JAX/Flax-NNX port of neuralforecast.VanillaTransformer). A classic encoder-decoder Transformer (Informer baseline). The encoder embeds the lookback window (circular token conv + fixed sinusoidal positional embedding) and runs full softmax self-attention; the decoder embeds concat(last label_len of input, zeros for h) and runs full self-attention + cross-attention against the encoder output, then projects to one channel and returns the last h steps. NF default scaler_type="identity" is mirrored (trained in raw 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).

__init__(self, h, input_size= -1, hidden_size= 128, n_heads= 4, conv_hidden_size= 32, encoder_layers= 2, decoder_layers= 1, decoder_input_size_multiplier= 0.5, dropout= 0.05, use_boxcox= False, activation= 'gelu', max_steps= 5000, learning_rate= 0.0001, windows_batch_size= 1024, random_seed= 1, alias= 'VanillaTransformer', loss= 'mae')

Initialize a VanillaTransformer forecaster. Stores hyperparameters; the network is built lazily at fit time. Defaults match neuralforecast.VanillaTransformer (hidden_size=128, n_head=4, conv_hidden_size=32, encoder_layers=2, decoder_layers=1, decoder_input_size_multiplier=0.5, dropout=0.05, activation="gelu", max_steps=5000, learning_rate=1e-4, windows_batch_size=1024, scaler_type="identity"). 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. use_boxcox (default False) applies a variance-stabilizing Box-Cox transform before modelling and inverts it on the forecast, mirroring :class:chronax.models.TBATS. It requires strictly positive values. Left off, the model is a faithful port of neuralforecast.VanillaTransformer.

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-1	Resolves to 3 * h if -1.
hidden_size	int	128	(undocumented)
n_heads	int	4	(undocumented)
conv_hidden_size	int	32	(undocumented)
encoder_layers	int	2	(undocumented)
decoder_layers	int	1	(undocumented)
decoder_input_size_multiplier	float	0.5	(undocumented)
dropout	float	0.05	(undocumented)
use_boxcox	bool	False	Applies a variance-stabilizing Box-Cox transform before modelling and inverts it on the forecast. Requires strictly positive values.
activation	str	"gelu"	"gelu" or "relu".
max_steps	int	5000	(undocumented)
learning_rate	Union[float, Callable[[int], float]]	1e-4	A scalar or an optax.ScalarOrSchedule. If callable, must be picklable.
windows_batch_size	int	1024	(undocumented)
random_seed	int	1	(undocumented)
alias	str	"VanillaTransformer"	(undocumented)
loss	Union[str, LossFn]	"mae"	A registry name ("mae"/"mse"/"huber") or a callable. If callable, must be picklable.

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

Fit the network on a 1-D series of length >= input_size + h.

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
X	jnp.ndarray \| None	None	(undocumented)

Returns: Self (the fitted forecaster; sets self.model_). Raises: NotImplementedError (If X is provided). ValueError (If y is not 1-D or too short). ValueError (If use_boxcox=True requires strictly positive series values).

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

Forecast h steps (1 <= h <= self.h) from the fitted context.

Parameter	Type	Default	Description
h	int	-	(undocumented)
X	jnp.ndarray \| None	None	(undocumented)
level	list[int \| float] \| None	None	(undocumented)

Returns: dict Keys: {"mean": jnp.ndarray}. Includes prediction interval keys (e.g., "lower_90", "upper_90") if level is provided and conformal_params is set. Raises: RuntimeError (If fit(y) has not been called). ValueError (If h is non-positive or greater than self.h). ValueError (If level is provided but model.conformal_params is not set).

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

Stateless fit-then-predict on y. Optionally add "fitted".

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
X	jnp.ndarray \| None	None	(undocumented)
X_future	jnp.ndarray \| None	None	(undocumented)
level	list[int \| float] \| None	None	(undocumented)
fitted	bool	False	(undocumented)

Returns: dict Keys: {"mean": jnp.ndarray}. Includes prediction interval keys if level is provided, and "fitted": jnp.ndarray if fitted=True. Raises: NotImplementedError (If X or X_future is provided).