predictive.predict_in_sample()

Sample the in-sample posterior predictive of the obs site.

Usage

Source

predictive.predict_in_sample(
    rng_key,
    model,
    posterior,
    covariates,
    *,
    batch_size=None,
    parallel=True,
    device=None
)

Runs Predictive with the in-sample covariates and the supplied posterior latent draws. Unlike forecast() there is no forecast horizon: covariates span only the observed window, so the model’s obs site is sampled at every step. The number of predictive samples equals the leading (sample) axis of posterior (see draw_posterior()). The model gets data=None, so the draws have the axes that the model gives obs, and the driver has no data to check them against: a predictor shared by n series without a size-n plate at dim=-1 draws one column. forecast() and to_datatree() get data and raise ValueError in that case.

Parameters

rng_key: Array

PRNG key.

model: ForecastModel

The forecasting model callable (the same one that produced posterior).

posterior: Mapping[str, ArrayLike]

Posterior samples of the latent sites of the training run, sample axis leading (mcmc.get_samples() of the fit, or draw_posterior()). The output of a device="host" stage (CPU-committed jax leaves or NumPy leaves) is accepted directly. Forecast-horizon keys follow the rule of forecast(): one with values raises ValueError, because it marks the output of a numpyro.infer.Predictive call over a horizon, and one with an empty leaf is dropped. An "obs" key is dropped too, so the draws come from the latent sites: Predictive substitutes every posterior sample site by name, and an "obs" key (the output of a Predictive call with data, or of a prior predictive call) would make every draw equal to it. Only "obs" is dropped: keep the other sites that the model observes when it gets data (a second likelihood) out of posterior, or they are substituted.

covariates: Array

Covariates with time at axis -2 spanning the observed window. Its time length must match the data the posterior was fit on, since the in-sample latent sites are sized to that window. An innovations() site of another length raises ValueError however the model slices its "obs" site. For other latents the driver checks only that the "obs" draws have one step per row, and a markov_series() latent replays a shorter posterior as a prefix and draws the remaining steps from the prior (a longer one raises NumPyro’s RuntimeError).

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). See forecast() for how this interacts with batch_size.

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; the same placement contract as the device argument of draw_posterior() ("host" for pageable host memory as a CPU-committed jax.Array, or a NumPy array when no CPU backend is initialized; "numpy"; "pinned_host"; a jax.Device or platform name; None for the default device), including its mixing rules for host-committed results. With batch_size set on an accelerator, any host target bounds accelerator memory by a single chunk instead of the full (sample, time, 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. The result feeds straight into to_datatree() and the batch_size-chunked evaluation metrics in evaluate, which accept host-resident draws.

Returns

Num[Array, " sample *batch time obs"]

In-sample posterior-predictive draws of the obs site (with device="host" committed to the CPU device, or a NumPy array when no CPU backend is initialized).

Raises

ValueError

If posterior is empty, a posterior leaf has no sample axis, or the leaves disagree on the sample-axis length, if a forecast-horizon key of posterior (one that ends in "_future" or contains "_future_", or "forecast") has a non-empty leaf, if "obs" is the only key of posterior, if covariates has fewer than two axes (time must be at axis -2) or no time steps, if a posterior value of an innovations() site does not have one step per covariates row (a posterior drawn for a different training length), or if the "obs" draws do not have time at axis -2 with one step per covariates row.

HostMemoryKindError

If device="pinned_host" is requested on a device that exposes no host memory kind (see _host_memory_kind()).

DevicePlatformError

If device names a platform whose backend is not initialized (see resolve_device()).

Warns

UserWarning

If device="cpu" is requested and the JAX CPU backend is not initialized, so the draws take the NumPy path of "host" instead.

Notes

Like forecast(), each call compiles the predictive once (shared by every batch_size chunk) and keeps no compiled program afterwards, so the draws reflect the model as it is at call time and model need not be hashable.