build_windows

patchtst_training.build_windows

Return [n_windows, input_size+h] rolling windows; step=1.

build_windows(y, input_size, h)

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
input_size	int	-	(undocumented)
h	int	-	(undocumented)

Returns: jnp.ndarray

forward_loss

patchtst_training.forward_loss

Forward + point loss in ORIGINAL scale (RevIN denorms inside the net).

forward_loss(model, windows, *, h, input_size, loss_fn=mae)

Parameter	Type	Default	Description
model	PatchTSTNet	-	(undocumented)
windows	jnp.ndarray	-	[B, input_size+h] -> scalar.
h	int	-	(undocumented)
input_size	int	-	(undocumented)
loss_fn	LossFn	mae	(undocumented)

Returns: jnp.ndarray

train

patchtst_training.train

Train model in place via a single nnx.scan. Returns per-step losses.

The whole loop is one nnx.scan (carry = (model, optimizer)), which keeps the function jax.vmap-traceable for BaseForecaster.conformity_scores. Window sampling replicates neuralforecast's REGIME-DEPENDENT scheme (_base_model.py training_step): when n_windows < windows_batch_size NF draws windows_batch_size indices WITH replacement (oversampling with duplicates — the regime every small benchmark series hits, e.g. ~24 windows for AirlinePassengers or ~245 for DailyFemaleBirths, both << 1024); otherwise it takes a without-replacement permutation of windows_batch_size windows. Getting this branch right is load-bearing for accuracy parity, so we do NOT collapse it to full-batch. Note: batches materializes a [max_steps, windows_batch_size, input_size+h] tensor up front (mirrors the GRU sibling's pre-sampling). At the benchmark defaults that is ~1.9 GB resident — acceptable for an offline benchmark, but reduce windows_batch_size or max_steps if memory-constrained.

train(model, y, *, h, input_size, max_steps, windows_batch_size, lr, seed, loss_fn=mae)

Parameter	Type	Default	Description
model	PatchTSTNet	-	(undocumented)
y	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
input_size	int	-	(undocumented)
max_steps	int	-	(undocumented)
windows_batch_size	int	-	(undocumented)
lr	optax.ScalarOrSchedule	-	(undocumented)
seed	int	-	(undocumented)
loss_fn	LossFn	mae	(undocumented)

Returns: jnp.ndarray

predict_step

patchtst_training.predict_step

Forecast next h steps from the final input_size of y. Returns (h,).

predict_step(model, y, *, h, input_size)

Parameter	Type	Default	Description
model	PatchTSTNet	-	(undocumented)
y	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
input_size	int	-	(undocumented)

Returns: jnp.ndarray