Window construction and JIT/scan training for MLP.

build_windows

def build_windows(y: jnp.ndarray, input_size: int, h: int) -> tuple[jnp.ndarray, jnp.ndarray]

Rolling windows over y, right-padded with h zeros.

Padding keeps every window with at least one real target and lets the newest observations appear as training contexts; the padded tail is masked out of the loss (matters on trending series, where the most recent regime is the forecast-relevant one). Returns (windows [n, input_size+h], target_mask [n, h]) with n = len(y) - input_size; the mask is 1.0 at real target positions and 0.0 in the zero-padded tail.

build_exog_windows

def build_exog_windows(arr: jnp.ndarray, input_size: int, h: int, n_windows: int, span: str) -> jnp.ndarray

Rolling windows of an exog array [T, F], right-padded with h zero rows to match build_windows (late windows see zeros in the padded tail of each exog channel).

span="input" -> [n, input_size, F] (encoder window); span="full" -> [n, input_size+h, F] (future-known spanning input + horizon).

train_on_windows

def train_on_windows(net, y_windows, target_mask, *, h, input_size, max_steps, windows_batch_size, lr, seed, loss_fn, scaler, futr_windows=None)

Train net in place on prebuilt window arrays via one nnx.scan.

Accepts pooled windows from any number of series (cross-learning); batch sampling is uniform over the pooled set — with replacement when there are fewer windows than windows_batch_size, without replacement otherwise. Returns per-step losses.

train

def train(net, y, *, h, input_size, max_steps, windows_batch_size, lr, seed, loss_fn, scaler, futr_exog=None)

Build windows from a single series and train (see train_on_windows).

predict_step

def predict_step(net, y_context, *, h, input_size, scaler, futr_full=None)

Forecast next h steps from per-series contexts, in the original scale (point/quantile heads only — distribution heads go through predict_params).

y_context is [L] (one series) or [B, L] (a batch of series tails); returns [h, multiplier] / [B, h, multiplier] accordingly. futr_full ([input_size+h, F], history + horizon) is shared across a batch of contexts.

predict_params

def predict_params(net, y_context, *, input_size, scaler, loss_fn, futr_full=None)

Distribution parameters for the next h steps, in the ORIGINAL scale.

y_context is [L] (one series) or [B, L] (a batch of contexts — per-series tails); returns the loss's decoupled parameter tuple with arrays [h, K] / [B, h, K] accordingly. futr_full ([input_size+h, F], history + horizon) is shared across a batch of contexts.