deepnpts_module

Flax NNX module for the DeepNPTS forecaster.

DeepNPTSNet

chronax.models.deepnpts_module.DeepNPTSNet

Univariate DeepNPTS backbone: [B, input_size, 1] -> [B, h, 1].

No-exogenous case: the MLP input dimension equals input_size. n_layers blocks of Linear -> ReLU -> [BatchNorm] -> [Dropout] feed a final linear that produces input_size * h logits; these are reshaped to [B, L, h], softmaxed over L (the window axis), and used to weight-sum the raw in-sample values into the horizon. Identity scaling (NF default): the network operates in raw scale. During training pass deterministic=False; the nnx.Rngs dropout stream supplies noise and BatchNorm updates running stats.

__init__(self, *, h: int, input_size: int, hidden_size: int, n_layers: int, dropout: float, batch_norm: bool, rngs: nnx.Rngs)

Initializes the DeepNPTS network components (MLP layers, BatchNorm, Dropout, and output projection).

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
hidden_size	int	-	(undocumented)
n_layers	int	-	(undocumented)
dropout	float	-	(undocumented)
batch_norm	bool	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray

Performs the forward pass of the DeepNPTS network.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input time series data, shape [B, L, 1].
deterministic	bool	-	If True, disables dropout and uses running averages for BatchNorm.

Returns: jnp.ndarray (The forecast, shape [B, h, 1]).