chronax.data

Data pipeline for DeepAR model (JAX/FLAX, NumPy-free implementation). The pipeline produces jnp.ndarray batches compatible with the DeepAR model. It maintains two masks following Nixtla's vocabulary: - available_mask: 1 where historic target is real, 0 where padded - sample_mask: 1 on horizon steps included in loss, 0 otherwise (for holidays/outliers)

pad_sequence

chronax.data.pad_sequence

Left-pad / left-truncate a 1D series to target_len (JAX version).

pad_sequence(y, target_len, pad_value=0.0)

Parameter	Type	Default	Description
y	jnp.ndarray	-	1D array of values
target_len	int	-	Desired output length
pad_value	float	0.0	Value to use for padding

Returns: Tuple[jnp.ndarray, jnp.ndarray] ((padded_array, mask) where mask is 1 on real data, 0 on padding).

align_covariates

chronax.data.align_covariates

Align exogenous arrays to model expectations (JAX version).

align_covariates(hist_exog, futr_exog, stat_exog, input_size, h)

Parameter	Type	Default	Description
hist_exog	Optional[jnp.ndarray]	-	Historic exogenous features, shape [T, X] or None
futr_exog	Optional[jnp.ndarray]	-	Future exogenous features, shape [T+h, F] or None
stat_exog	Optional[jnp.ndarray]	-	Static exogenous features, shape [S] or None
input_size	int	-	History window length
h	int	-	Forecast horizon

Returns: Dict[str, Optional[jnp.ndarray]] (Dict with aligned JAX arrays).

create_batch

chronax.data.create_batch

Build a single JAX batch from a list of time series.

create_batch(y_series, input_size, h, hist_exog_list=None, futr_exog_list=None, stat_exog_list=None, sample_mask_list=None)

Parameter	Type	Default	Description
y_series	List[jnp.ndarray]	-	List of 1D JAX arrays, each of length T_i
input_size	int	-	History window length L
h	int	-	Forecast horizon
hist_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Per-series [T_i, X] historic exog, or None
futr_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Per-series [T_i + h, F] future exog, or None
stat_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Per-series [S] static exog, or None
sample_mask_list	Optional[List[Optional[jnp.ndarray]]]	None	Per-series [h] horizon masks, or None

Returns: Dict[str, Optional[jnp.ndarray]] (Dict with JAX arrays ready for model).

batch_generator

chronax.data.batch_generator

Iterate over y_series yielding fully-prepared JAX batches.

batch_generator(y_series, input_size, h, batch_size, hist_exog_list=None, futr_exog_list=None, stat_exog_list=None, sample_mask_list=None, shuffle=True, seed=42)

Parameter	Type	Default	Description
y_series	List[jnp.ndarray]	-	List of 1D JAX arrays
input_size	int	-	History window length
h	int	-	Forecast horizon
batch_size	int	-	Batch size
hist_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Historic exogenous features per series
futr_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Future exogenous features per series
stat_exog_list	Optional[List[Optional[jnp.ndarray]]]	None	Static exogenous features per series
sample_mask_list	Optional[List[Optional[jnp.ndarray]]]	None	Sample masks per series
shuffle	bool	True	Whether to shuffle series
seed	int	42	Random seed for shuffling

Yields: Iterator[Dict[str, Optional[jnp.ndarray]]] (Dict with batched JAX arrays).