NBEATSConfig

chronax.model.NBEATSConfig

Hyperparameters for :class:NBEATS.

__init__(self, h=24, input_size=48, stack_types=('identity', 'trend', 'seasonality'), n_blocks=(1, 1, 1), mlp_units=512, mlp_layers=4, n_harmonics=2, n_basis=2, basis='polynomial', activation='relu', shared_weights=False, dropout_prob=0.0, layer_norm=True)

Parameter	Type	Default	Description
h	int	24	Forecast horizon.
input_size	int	48	History window length L fed to each block.
stack_types	Tuple[str, ...]	('identity', 'trend', 'seasonality')	Sequence of block-stack types; each element is one of "identity", "trend", or "seasonality".
n_blocks	Tuple[int, ...]	(1, 1, 1)	Number of blocks per stack (must have same length as stack_types).
mlp_units	int	512	Hidden dimension used in every Dense layer of every block.
mlp_layers	int	4	Number of hidden Dense layers per block (output layer is extra). Matches len(mlp_units) in NF when all units are equal.
n_harmonics	int	2	Fourier harmonics for the seasonality basis (n_harmonics=2 gives the NF default).
n_basis	int	2	Polynomial degree for the trend basis; basis size = n_basis+1.
basis	str	'polynomial'	Trend basis type — "polynomial" matches NF's TrendBasis with basis="polynomial". Other choices: "legendre", "chebyshev", "changepoint".
activation	str	'relu'	Activation function name ("relu" matches NF default).
shared_weights	bool	False	When True all blocks within a stack share parameters.
dropout_prob	float	0.0	Dropout rate (0 = disabled, matches NF default).
layer_norm	bool	True	When True insert LayerNorm after every hidden Dense.

NBEATSBlock

chronax.model.NBEATSBlock

Single N-BEATS block: FC stack → theta → basis projection → (backcast, forecast).

__init__(self, config, block_type)

Parameter	Type	Default	Description
config	NBEATSConfig	-	Shared model config (provides mlp_units, mlp_layers, etc.).
block_type	str	-	One of "identity", "trend", "seasonality".

__call__(self, residuals, deterministic=True) -> Tuple[jnp.ndarray, jnp.ndarray]

Parameter	Type	Default	Description
residuals	jnp.ndarray	-	[B, L]
deterministic	bool	True	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

NBEATS

chronax.model.NBEATS

N-BEATS with doubly-residual stacking.

Input insample_y should already be scaled (the forecaster applies RobustScaler per window before calling model.apply).

Forward pass:

1. Level init: forecast = broadcast(insample_y[:, -1], h) (Naive1 level — mirrors NeuralForecast exactly).
2. residuals = flip(insample_y) (most-recent step first, NF convention).
3. For each block: forecast += block_forecast; residuals -= backcast.
4. Return forecast.

__init__(self, config)

Parameter	Type	Default	Description
config	NBEATSConfig	-	:class:NBEATSConfig specifying architecture.

__call__(self, insample_y, insample_mask=None, deterministic=True) -> jnp.ndarray

Parameter	Type	Default	Description
insample_y	jnp.ndarray	-	[B, L]
insample_mask	Optional[jnp.ndarray]	None	[B, L] 1=valid 0=padded
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray ([B, h])