chronax.data

Functional data pipeline for the Chronax TSMixer model. TSMixer is a multivariate model — all N series are processed jointly in a single [B, L, N] window. The pipeline is therefore organised around a single 2D array [T, N] rather than a list of univariate series as in the RNN module.

create_windows(y, input_size, h)

chronax.data.create_windows

Extract all valid sliding windows from a [T, N] multivariate series.

If the series is shorter than input_size + h, the history is zero-padded on the left and a single window is returned.

Parameter	Type	Default	Description
y	jnp.ndarray	-	[T, N] float array (or 1-D for univariate).
input_size	int	-	history window length L.
h	int	-	forecast horizon.

Returns: Tuple[jnp.ndarray, jnp.ndarray] * insample: [W, L, N] float32 — history windows. * outsample: [W, h, N] float32 — corresponding horizon targets.

make_batch(insample, outsample, indices)

chronax.data.make_batch

Build a JAX batch dict from pre-extracted window arrays.

This is a thin indexing wrapper intended for use inside the training loop where the full window cache has been transferred to device once and each step only selects a random subset.

Parameter	Type	Default	Description
insample	jnp.ndarray	-	[W, L, N] all history windows (device array).
outsample	jnp.ndarray	-	[W, h, N] all horizon windows.
indices	jnp.ndarray	-	[B] integer indices selecting windows for this batch.

Returns: Dict[str, jnp.ndarray] * A dictionary: {"insample_y": [B, L, N], "outsample_y": [B, h, N], "sample_mask": [B, h, N]} with all-ones mask.