functional.prediction.forecast()

Sample forecasts for the steps in [t, duration) from a posterior.

Usage

functional.prediction.forecast(
    rng_key,
    model,
    posterior,
    data,
    covariates,
    *,
    batch_size=None,
    parallel=True,
    device=None
)

Runs Predictive with full-horizon covariates and the in-sample data: the in-sample latent sites are drawn from posterior while the _future suffix is drawn from the prior, and the "forecast" site is returned. The number of forecast samples equals the leading (sample) axis of posterior (see ~numpyro_forecast.functional.posterior.draw_posterior()).

Parameters

rng_key: Array: PRNG key.
model: ForecastModel: The forecasting model callable (the same one that produced posterior).
posterior: Mapping[str, Array | np.ndarray]: Posterior samples of the latent sites, sample axis leading.
data: Array: Observed data with time at axis -2 and length t.
covariates: Array: Covariates with time at axis -2 and length duration > t.
batch_size: int | None = None: Optional chunk size for sampling (caps peak memory).
parallel: bool = True: Whether Predictive vectorizes over the sample axis with vmap (True, faster, higher peak memory) or maps it serially with lax.map (False). With parallel=True the samples in each batch_size chunk are vectorized while the chunks are looped over, so batch_size remains the peak-memory governor. The two settings produce the same draws up to floating-point reduction order.
device: jax.Device | str | None = None: Where each chunk of draws is placed as soon as it is drawn and where the stitched result lives. "host" copies every chunk to host memory with jax.device_get() and returns a NumPy array; it needs no CPU backend, so it works even when numpyro.set_platform("cuda") (or jax_platforms) leaves only an accelerator backend initialized, which makes it the recommended choice on GPU. A jax.Device or platform name like "cpu" commits the draws to that device instead ("cpu" falls back to "host" with a UserWarning when the CPU backend is not initialized). With batch_size set on an accelerator, either bounds accelerator memory by a single chunk instead of the full (sample, future, obs) array; the draw values are unchanged, only where the result lives. The bound requires batch_size strictly below the sample count: at or above it, the single-shot path runs and the full array is materialized on the default device before the one transfer. None keeps everything on the default device.

Returns

Num[Array, " sample *batch future obs"] | Num[np.ndarray, " sample *batch future obs"]: Forecast samples over the future = duration - t horizon (floating point for continuous observations, integer for discrete/count models built with ~numpyro_forecast.functional.models.predict_glm(); a NumPy array when device resolves to "host").

Raises

ValueError: If covariates does not extend beyond data along the time axis.

Notes

Chunking is a memory knob, not a reproducibility knob: reproducibility is per (rng_key, batch_size). Every chunk shares the exact batch_size shape (the final chunk wraps around to re-used draws that are discarded), so the underlying _predict compiles exactly once for a fixed shape, but changing batch_size changes the PRNG stream layout and therefore the exact draws. device never changes the draws, only where they live.