---------------------------------------------------------------------- This is the API documentation for the numpyro_forecast library. ---------------------------------------------------------------------- ## Model building blocks Plain model functions that register the train/forecast sites for you. Horizon(data: jax.Array | None, t_obs: int, future: int, duration: int) -> None The train/forecast split for a single model call. An immutable value derived once per model call from the covariate and data shapes by `from_data()`; every building block (`innovations()`, `markov_series()`, `ssoe()`, `predict()`) takes it as its first argument. A JAX pytree (an `equinox.Module`): ``data`` is the only leaf, while ``t_obs``, ``future`` and ``duration`` are static metadata. A jitted function that takes a `Horizon` can therefore use the three integers as shapes and recompiles once per horizon length. Attributes ---------- data Observed in-sample data with time at axis ``-2`` (``None`` during pure prior sampling). t_obs Number of observed (in-sample) time steps ``t``. future Number of forecast time steps ``f`` (``0`` while training). duration Total horizon length ``t + future`` (in time steps). Transition Type alias. Type aliases are created through the type statement:: type Alias = int In this example, Alias and int will be treated equivalently by static type checkers. At runtime, Alias is an instance of TypeAliasType. The __name__ attribute holds the name of the type alias. The value of the type alias is stored in the __value__ attribute. It is evaluated lazily, so the value is computed only if the attribute is accessed. Type aliases can also be generic:: type ListOrSet[T] = list[T] | set[T] In this case, the type parameters of the alias are stored in the __type_params__ attribute. See PEP 695 for more information. Advance Type alias. Type aliases are created through the type statement:: type Alias = int In this example, Alias and int will be treated equivalently by static type checkers. At runtime, Alias is an instance of TypeAliasType. The __name__ attribute holds the name of the type alias. The value of the type alias is stored in the __value__ attribute. It is evaluated lazily, so the value is computed only if the attribute is accessed. Type aliases can also be generic:: type ListOrSet[T] = list[T] | set[T] In this case, the type parameters of the alias are stored in the __type_params__ attribute. See PEP 695 for more information. innovations(h: numpyro_forecast.models.Horizon, name: str, prior: numpyro.distributions.distribution.Distribution, *, reparam: numpyro.infer.reparam.Reparam | None = None) -> jax.Array Sample conditionally iid per-step innovations over the full horizon. The in-sample portion is sampled under ``plate("time", t)`` with the fixed site ``name``; when forecasting, the horizon portion is sampled under a separate site ``f"{name}_future"`` and concatenated. The separate site keeps the guide shape fixed and lets ``Predictive`` draw the forecast suffix from the prior. Build the series arithmetically from the result (a random walk is ``jnp.cumsum(drift, axis=-2)``); a latent whose per-step distribution depends on the previous state is `markov_series()`, and a deterministic error-feedback recursion driven by the observed series is `ssoe()`. Parameters ---------- h The horizon for the current model call (see `Horizon`). name Base sample-site name for the in-sample latent. prior The per-step prior distribution, shared by the in-sample and forecast sites (each time plate expands a copy; the instance is never mutated). Its batch shape is the per-step shape, for example ``()`` for a scalar latent or ``(n_series,)`` under an enclosing series plate; the time axis comes from the plate. reparam Optional reparameterization (e.g. ``LocScaleReparam``) applied to both the in-sample and forecast sites. Returns ------- Array The latent over the full horizon with time at axis ``-2``. markov_series(h: numpyro_forecast.models.Horizon, name: str, init_carry: Carry, transition: Transition[Carry], xs: jaxtyping.PyTree[jax.Array] | None = None, *, advance: Advance[Carry] | None = None, plates: collections.abc.Sequence[tuple[str, int]] = (), reparam_config: collections.abc.Mapping[str, numpyro.infer.reparam.Reparam] | None = None) -> jax.Array Sample a Markov (state-space) latent over the full horizon. In-sample steps run in a ``numpyro.contrib.control_flow.scan`` with site ``name``; when forecasting, horizon steps run in a second scan with site ``f"{name}_future"`` **seeded by the final in-sample carry**. The guide never sees the future site (same invariant as `innovations()`), and under posterior replay the carry is a deterministic function of the replayed draws, so the forecast is conditioned through the state. Parameters ---------- h The horizon for the current model call (see `Horizon`). name Base sample-site name for the in-sample latent scan. init_carry Initial carry passed to the first transition. transition Per-step ``(carry, x_t) -> dist_t`` callable returning the distribution of the next latent (see `Transition`); the wrapper owns the ``numpyro.sample`` statement. xs Optional exogenous inputs over the full horizon: a PyTree of arrays with time at axis ``-2`` (a single array, a tuple, a dict, ...), moved leaf by leaf into scan layout internally; ``None`` for autonomous dynamics. advance Optional ``(carry, z_t, x_t) -> carry`` (see `Advance`) that builds the next carry from the sampled latent; ``None`` means the carry *is* the latent, ``carry_{t+1} = z_t``, so ``init_carry`` must be a single array shaped like one draw. A vector autoregression with ``p`` lags keeps a ``(p, obs)`` window here. plates ``(name, size)`` pairs opened **inside** the scan body around the sample statement (the only placement NumPyro supports for scan + plate). reparam_config Site-name -> `numpyro.infer.reparam.Reparam` mapping applied **inside** the scan body. Returns ------- Array The latent over the full horizon in package layout ``(*plate_batch, duration, obs)``. Raises ------ ValueError If forecasting without observed data (only reachable with a hand-built `Horizon`: `Horizon.from_data()` never sets ``future > 0`` without data), if the per-step shape lacks the observation dimension, if an enclosing plate is detected, or if ``advance`` is omitted and ``init_carry`` is not a single array with the shape and dtype of a draw. ssoe(h: numpyro_forecast.models.Horizon, name: str, y: jax.Array | None, init_carry: Carry, mean: SSOEMean[Carry], update: SSOEUpdate[Carry], noise_dist: numpyro.distributions.distribution.Distribution, xs: jaxtyping.PyTree[jax.Array] | None = None) -> numpyro_forecast.models.SSOEResult Run a single-source-of-error recursion over the full horizon. The building block for innovations state-space models (ARMA, exponential smoothing, Croston/TSB levels, censored autoregressions): a deterministic filter whose state is driven by the one-step-ahead *error* ``eps_t = y_t - mu_t``. In-sample it runs ``mean`` and ``update`` in a raw ``jax.lax.scan`` over the observed series ``y`` (no sample sites inside); when forecasting it draws iid future errors at the site ``f"{name}_future"`` from ``noise_dist`` under a ``plate("time_future", h.future)`` and runs a second scan from the final in-sample carry with ``y_t = mu_t + eps_t`` fed back through ``update``. The guide never sees the future site, because fitting always happens with ``future == 0``. Linear-Gaussian members (ARMA, additive exponential smoothing) can be marginalized exactly by a Kalman filter; the error-feedback form is the one that also covers the nonlinear members. The block registers nothing but the error site. The caller writes the likelihood against ``r.mu`` and registers ``numpyro.deterministic("forecast", r.y_future)`` when ``h.future > 0`` (an unconditional registration is harmless to `~~numpyro_forecast.predictive.forecast()`, but lands a size-0 variable in every posterior). Driver contract: `~~numpyro_forecast.predictive.predict_in_sample()` and `~~numpyro_forecast.convert.to_datatree()` call the model with ``data=None`` and read ``"obs"``, so ``y`` must come from ``covariates`` or be computed in the model, never from ``h.data``; `~~numpyro_forecast.predictive.forecast()` reads ``"forecast"``. **Frozen gates.** Route an update gate (Croston's demand indicator, an availability mask) through an ``xs`` leaf frozen over the horizon with `~~numpyro_forecast.arrays.pad_future()`, and read it from ``x_t``, never from ``y_t`` (over the horizon ``y_t = mu_t + eps_t`` is nonzero). With the gate off, ``update`` returns the carry unchanged and the forecast is the last level plus iid errors. `~~numpyro_forecast.evaluate.backtest()` and `~~numpyro_forecast.predictive.forecast()` hand the model *real* future covariate rows, so a gate sliced from the full covariates keeps updating on sampled values and leaks the test window; scenario inputs such as a future availability mask are the only thing to read from those rows. **Shapes.** Rows are ``(*batch, obs)``: a scalar state emits ``mu[None]`` and starts from ``init[None]`` (the block refuses a scalar or a wider mean because either would silently broadcast the likelihood into a ``(t, t)`` log-prob); a tuple carry with scalar leaves reads ``eps_t[0]`` (the ETS idiom). A panel puts the series on the observation axis: ``y`` is ``(t_obs, series)``, the carry ``(series,)``, and a ``noise`` sampled under ``plate("series")`` has exactly the batch shape ``(series,)`` the block needs. Batch dims to the left of time (``(B, t_obs, obs)``) take a ``(B, obs)`` carry and a ``(B, 1, obs)`` noise batch. Errors correlated across the observation axis are a multivariate ``noise_dist`` with event shape ``(obs,)``, which is how a vector autoregression enters the block (see `~~numpyro_forecast.var.var_step()`). Inputs are jax Arrays (the import hook rejects NumPy). With ``obs == 1`` a noise batch shape ``(future, 1)`` is indistinguishable from time and is consumed as such: per-step error scales, if that is what you meant. **Composition.** Two channels are two calls sharing the same ``Horizon``; each opens its own ``time_future`` plate. Scoping a helper that contains the call is fine (``handlers.scope`` prefixes the error site and the plate; use ``name="eps"`` inside a scoped helper so the site reads ``z_eps_future``); register ``"obs"`` and ``"forecast"`` outside any scope and build the forecast from the channels' ``y_future``. Parameters ---------- h The horizon for the current model call (see `Horizon`). name Base name of the error site; the future errors are drawn at ``f"{name}_future"``. y The driving series over the observed window, shape ``(*batch, t_obs, obs)`` with time at axis ``-2`` (integer counts are fine as long as the carry, hence the mean, is floating; the error promotes). Sliced from ``covariates`` or computed in the model; ``None`` (the value of ``h.data`` under ``data=None``) raises. init_carry Initial carry, any PyTree, already broadcast to the ``(*batch, obs)`` rows: a scalar level is ``init[None]``, a panel level ``(series,)``. Every leaf must keep its shape and dtype through ``update``. mean ``(carry, x_t) -> mu_t`` (see `SSOEMean`): the mean for the current row (shape ``(*batch, obs)``, so a scalar state emits ``mu[None]``). update ``(carry, y_t, eps_t, x_t) -> carry`` (see `SSOEUpdate`): the next carry from the row's value and error. Over the horizon it receives the drawn ``eps_t`` (not a recomputed ``y_t - mu_t``, which can differ by an ulp); when the update needs the mean, call ``mean(carry, x_t)`` inside it. noise_dist Zero-centered per-step error distribution, either an elementwise family (event rank 0) or a multivariate family over the observation axis (event rank 1, e.g. ``dist.MultivariateNormal(jnp.zeros(obs), scale_tril=L)`` for shocks correlated across series, the VAR case). For a ``(t_obs, obs)`` series the batch shape is ``(obs,)`` (``()`` is fine when ``obs == 1``) for the elementwise form and ``()`` for the multivariate form; a batched ``(B, t_obs, obs)`` series takes ``(B, 1, obs)`` and ``(B, 1)`` respectively. Either way the draw under the time plate (``dim=-2`` for event rank 0, ``dim=-1`` for event rank 1) is exactly ``(*batch, future, obs)`` with the dtype of the means; event rank 2 or higher is rejected. xs Optional exogenous inputs over the full horizon: a PyTree of arrays with time at axis ``-2`` and ``duration`` rows (a single array, a tuple, a dict, ...), split at ``h.t_obs`` and handed to ``mean`` and ``update`` row by row as ``x_t``; ``None`` for autonomous dynamics. Returns ------- SSOEResult ``mu`` (in-sample means), ``mu_future`` and ``y_future`` (forecast means and sampled values; size-0 time axis while training). Raises ------ ValueError If ``y`` is ``None``, lacks the time or observation axis, or does not cover exactly ``h.t_obs`` rows; if an ``xs`` leaf lacks the axes or does not span ``h.duration`` rows; if ``mean`` returns a value without the observation axis or ``update`` a carry with a different tree structure, shape or dtype; if ``mean`` or ``update`` calls ``numpyro.sample``; if ``noise_dist`` has event rank 2 or higher, or has event rank 1 inside an enclosing plate at ``dim=-1``; or if it draws errors of the wrong shape or dtype. Examples -------- ARMA(1,1) (``y`` is the observed series routed through ``covariates``): >>> def mean(carry, _): ... y_prev, eps_prev = carry ... return mu + phi * y_prev + theta * eps_prev >>> def update(carry, y_t, eps_t, _): ... return y_t, eps_t >>> r = ssoe(h, "eps", y, (mu[None], jnp.zeros((1,))), mean, update, dist.Normal(0.0, sigma)) >>> numpyro.sample("obs", dist.Normal(r.mu, sigma), obs=h.data) >>> if h.future > 0: ... numpyro.deterministic("forecast", r.y_future) A gated level (Croston, TSB) with the gate frozen over the horizon: >>> def mean(level, _): ... return level >>> def update(level, y_t, _, gate_t): ... return jnp.where(gate_t, alpha * y_t + (1 - alpha) * level, level) >>> gate_full = pad_future(gate, h.future) >>> r = ssoe(h, "eps", y, init[None], mean, update, dist.Normal(0.0, noise), xs=gate_full) SSOEMean Type alias. Type aliases are created through the type statement:: type Alias = int In this example, Alias and int will be treated equivalently by static type checkers. At runtime, Alias is an instance of TypeAliasType. The __name__ attribute holds the name of the type alias. The value of the type alias is stored in the __value__ attribute. It is evaluated lazily, so the value is computed only if the attribute is accessed. Type aliases can also be generic:: type ListOrSet[T] = list[T] | set[T] In this case, the type parameters of the alias are stored in the __type_params__ attribute. See PEP 695 for more information. SSOEUpdate Type alias. Type aliases are created through the type statement:: type Alias = int In this example, Alias and int will be treated equivalently by static type checkers. At runtime, Alias is an instance of TypeAliasType. The __name__ attribute holds the name of the type alias. The value of the type alias is stored in the __value__ attribute. It is evaluated lazily, so the value is computed only if the attribute is accessed. Type aliases can also be generic:: type ListOrSet[T] = list[T] | set[T] In this case, the type parameters of the alias are stored in the __type_params__ attribute. See PEP 695 for more information. SSOEResult(mu: jaxtyping.Float[Array, '*batch t_obs obs'], mu_future: jaxtyping.Float[Array, '*batch future obs'], y_future: jaxtyping.Float[Array, '*batch future obs']) The means and sampled future values produced by `ssoe()`. A named tuple, hence a JAX pytree with three array leaves: it unpacks as ``mu, mu_future, y_future = ssoe(...)`` and passes through ``jax.tree.map``, ``jax.jit`` and ``jax.vmap`` unchanged. Attributes ---------- mu In-sample one-step-ahead means, shape ``(*batch, t_obs, obs)``: the predictor the caller writes its likelihood against. mu_future Forecast-horizon one-step-ahead means, shape ``(*batch, future, obs)`` (a size-0 time axis while training). y_future Sampled future values ``mu_future + eps``, shape ``(*batch, future, obs)`` (a size-0 time axis while training); the caller registers them as the ``"forecast"`` deterministic when ``h.future > 0``. predict(h: numpyro_forecast.models.Horizon, obs_dist: numpyro.distributions.distribution.Distribution | collections.abc.Callable[[jax.Array], numpyro.distributions.distribution.Distribution], prediction: jax.Array) -> None Register the observation and forecast sites for the model. ``prediction`` is the deterministic predictor over the full horizon (time at axis ``-2``, shape ``(*batch, duration, obs)``). ``obs_dist`` takes one of two forms. A `numpyro.distributions.Distribution` is zero-centered observation noise (e.g. ``dist.StudentT(nu, 0.0, sigma)``) shifted onto the predictor with `~~numpyro_forecast.surgery.shift_loc()`, which also owns the multivariate-normal layout check. A callable is a link mapping the predictor to the observation distribution directly (the GLM form, e.g. ``lambda eta: dist.Poisson(jnp.exp(eta))``). Either way, while training the observation site ``"obs"`` is observed; while forecasting the in-sample prefix is observed and the forecast suffix is sampled at ``"obs_future"`` and exposed as the ``"forecast"`` deterministic site that `~~numpyro_forecast.predictive.forecast()` reads. The observation distribution must support time-axis surgery (`~~numpyro_forecast.surgery.slice_time()` / `~~numpyro_forecast.surgery.prefix_condition()`), i.e. an elementwise family or a registered one. Parameters ---------- h The horizon for the current model call (see `Horizon`). obs_dist Zero-centered noise distribution (shifted onto ``prediction``) or a link callable from the full-horizon predictor to the observation distribution. prediction The deterministic predictor over the full horizon, time at axis ``-2``. Raises ------ TypeError If ``obs_dist`` is a callable that does not return a distribution (for example a link that returns the predictor itself). RuntimeError If forecasting (``future > 0``) but no observed data is available. ValueError If the observation distribution has discrete support but ``h.data`` is not integer-dtyped (the usual mistake for count models). PlateName(*values) Names of the plates the building blocks open, as ``str`` members. ``TIME`` is the in-sample plate of `innovations()`; ``TIME_FUTURE`` is the forecast-horizon plate of `innovations()` and `ssoe()`. `~~numpyro_forecast.reparam.time_reparam()` targets the sites under ``TIME`` by name, so both modules read this enum rather than spelling the literal. ## Vector autoregression VAR components that compose with `ssoe` and `markov_series`: conditional mean, step factory, companion matrix, impulse responses. var_mean(phi: jaxtyping.Float[Array, '*#batch lags obs obs'], lags: jaxtyping.Float[Array, '*#batch lags obs'], intercept: jaxtyping.Float[Array, '*#batch obs'] | None = None) -> jaxtyping.Float[Array, '*batch obs'] Compute the VAR conditional mean of the next row from a lag window. $$ \mu_t = c + \sum_{l=1}^{p} \Phi_l \, y_{t-l}, $$ with ``phi[..., l - 1, :, :]`` holding $\Phi_l$ and ``lags[..., -l, :]`` holding $y_{t-l}$ (most recent row last). Parameters ---------- phi Coefficient tensor ``(*batch, lags, obs, obs)``; see the module docstring for the index convention. lags Lag window ``(*batch, lags, obs)`` in natural time order. intercept Optional intercept $c$ of shape ``(*batch, obs)``; ``None`` for a zero-mean recursion (the impulse response case). Returns ------- Float[Array, "*batch obs"] The conditional mean row, broadcast over the batch axes of the inputs. Examples -------- A latent VAR under `~~numpyro_forecast.models.markov_series()`: the transition returns the next-row distribution and ``advance`` shifts the lag window, with ``phi`` and the Cholesky factor ``scale_tril`` sampled outside: ```python def transition(window, _): return dist.MultivariateNormal(var_mean(phi, window), scale_tril=scale_tril) def advance(window, z_t, _): return jnp.concatenate([window[..., 1:, :], z_t[..., None, :]], axis=-2) z = markov_series(h, "z", jnp.zeros((n_lags, n_obs)), transition, advance=advance) ``` var_step(phi: jaxtyping.Float[Array, '*#batch lags obs obs'], intercept: jaxtyping.Float[Array, '*#batch obs'] | None = None) -> tuple[SSOEMean[jaxtyping.Float[Array, '*batch lags obs']], SSOEUpdate[jaxtyping.Float[Array, '*batch lags obs']]] Build the `~~numpyro_forecast.models.ssoe()` mean and update of a VAR. The carry is the lag window ``(*batch, lags, obs)``. ``mean`` emits `var_mean()` of the window as the one-step-ahead mean and, given the row's value, ``update`` drops the oldest row and appends the new one. Both ignore their exogenous input ``x_t``; add regressors by wrapping ``mean`` (a VARX): ```python mean, update = var_step(phi, intercept) def mean_x(carry, x_t): return mean(carry, x_t) + beta @ x_t ``` The pair knows nothing about priors: ``phi`` and ``intercept`` are whatever the model sampled (a weakly informative ``Normal``, the moments of `~~numpyro_forecast.priors.minnesota_prior()`, a hierarchical prior, ...), so changing the prior never touches the recursion. Parameters ---------- phi Coefficient tensor ``(*batch, lags, obs, obs)``; see the module docstring. intercept Optional intercept of shape ``(*batch, obs)``. Returns ------- tuple[SSOEMean[Float[Array, "*batch lags obs"]], SSOEUpdate[Float[Array, "*batch lags obs"]]] A ``(mean, update)`` pair for `~~numpyro_forecast.models.ssoe()`: ``mean(carry, x_t)`` is `var_mean()` of the window and ``update(carry, y_t, eps_t, x_t)`` drops the oldest row and appends ``y_t``. Raises ------ ValueError At step time, if the carry does not hold exactly ``phi.shape[-3]`` rows (the usual cause is an ``init_carry`` with the wrong number of lags). Examples -------- An observed VAR(p) conditioned on its first ``p`` rows. ``y_init`` is the seed window, a constant closed over by the model; the likelihood rows travel through ``covariates`` (padded with `~~numpyro_forecast.arrays.pad_future()` to fix the forecast horizon) and through ``data``: ```python def var_model(covariates, data=None): h = Horizon.from_data(covariates, data) y = covariates[..., : h.t_obs, :] intercept = jnp.asarray( numpyro.sample("intercept", dist.Normal(0.0, 1.0).expand([k]).to_event(1)) ) sigma = jnp.asarray(numpyro.sample("sigma", dist.HalfNormal(1.0).expand([k]).to_event(1))) l_omega = jnp.asarray(numpyro.sample("l_omega", dist.LKJCholesky(k, concentration=1.0))) phi = jnp.asarray( numpyro.sample("phi", dist.Normal(0.0, 1.0).expand([p, k, k]).to_event(3)) ) scale_tril = sigma[..., :, None] * l_omega noise = dist.MultivariateNormal(jnp.zeros(k), scale_tril=scale_tril) mean, update = var_step(phi, intercept) r = ssoe(h, "eps", y, y_init, mean, update, noise) numpyro.sample("obs", dist.MultivariateNormal(r.mu, scale_tril=scale_tril), obs=h.data) if h.future > 0: numpyro.deterministic("forecast", r.y_future) ``` Closing over ``y_init`` is right for a single fit but wrong under `~~numpyro_forecast.evaluate.backtest()`, which slices ``covariates`` per window: for a backtest, ship the seed rows inside ``covariates`` and slice the carry from them in the model. companion_matrix(phi: jaxtyping.Float[Array, '*#batch lags obs obs']) -> jaxtyping.Float[Array, '*batch lags_obs lags_obs'] Stack the VAR coefficients into the companion form of a VAR(1). $$ F = \begin{bmatrix} \Phi_1 & \Phi_2 & \cdots & \Phi_{p-1} & \Phi_p \\ I & 0 & \cdots & 0 & 0 \\ 0 & I & \cdots & 0 & 0 \\ \vdots & & \ddots & & \vdots \\ 0 & 0 & \cdots & I & 0 \end{bmatrix}, $$ a square matrix of size ``lags * obs``. The companion state stacks the lag window **most recent first**, the reverse of the ``lags`` layout: ``s = lags[..., ::-1, :].reshape(*batch, lags * obs)``, and then ``(F @ s)[..., :obs]`` equals `var_mean()` without intercept. The VAR is stable exactly when every eigenvalue of $F$ has modulus below one; that is the condition for a finite unconditional mean $(I - \sum_l \Phi_l)^{-1} c$ and for the impulse responses of `impulse_response()` to die out. Parameters ---------- phi Coefficient tensor ``(*batch, lags, obs, obs)``. Returns ------- Float[Array, "*batch lags_obs lags_obs"] The companion matrix per batch element. Notes ----- JAX implements the eigenvalues of a nonsymmetric matrix on CPU only, so compute the spectral radius of posterior draws with NumPy: ``np.abs(np.linalg.eigvals(np.asarray(companion_matrix(phi)))).max(-1)``. impulse_response(phi: jaxtyping.Float[Array, '*#batch lags obs obs'], horizon: int, *, scale_tril: jaxtyping.Float[Array, '*#batch obs obs'] | None = None, cumulative: bool = False) -> jaxtyping.Float[Array, '*batch steps obs obs'] Compute the impulse responses (moving-average coefficients) of a VAR. Inverting the lag polynomial of $y_t = c + \sum_{l=1}^{p} \Phi_l y_{t-l} + \varepsilon_t$ gives the moving-average representation $y_t = \mu + \sum_{h \ge 0} \Psi_h \varepsilon_{t-h}$ whose coefficients follow the recursion $$ \Psi_0 = I, \qquad \Psi_h = \sum_{j=1}^{\min(h, p)} \Phi_j \Psi_{h-j} \quad (h \ge 1). $$ ``out[..., h, i, j]`` is the response of variable ``i``, ``h`` steps after a unit shock to variable ``j`` at step ``0`` (a unit increase of the reduced-form residual $\varepsilon_{t, j}$ with the other residuals held at zero). With ``scale_tril`` the responses are to *orthogonalized* shocks: for a factor $B$ with $\Sigma = B B^\top$ (the Cholesky factor $L$ of the shock covariance, so $\varepsilon_t = L u_t$ with $u_t$ standard normal) the function returns $\Theta_h = \Psi_h B$, the response to a one-standard- deviation shock $u_{t, j}$. With the Cholesky factor this is the recursive identification: the first series in the ordering responds to its own shock only at $h = 0$, the second to the first two, and so on, so the ordering of the series is part of the model. The recursion is defined for any ``phi``; the moving-average representation, and $\Psi_h \to 0$, need a stable VAR (see `companion_matrix()`). The left recursion above and the right recursion $\Psi_h = \sum_j \Psi_{h-j} \Phi_j$ agree, because both are the two-sided inverse of the lag polynomial. Parameters ---------- phi Coefficient tensor ``(*batch, lags, obs, obs)``. Posterior draws pass through the leading batch axis unchanged; no ``vmap`` is needed. horizon Largest step $H \ge 0$ to compute; the result holds ``horizon + 1`` steps including step ``0`` (like `~~numpyro_forecast.acf.acf()`, which includes lag ``0``). scale_tril Optional square factor ``(*batch, obs, obs)`` of the shock covariance (typically the Cholesky factor, ``sigma[..., :, None] * l_omega`` for an LKJ correlation factor). Its batch axes broadcast against those of ``phi``. cumulative If ``True``, return the running sum over steps, the response of the *level* when the modeled series are differences. Returns ------- Float[Array, "*batch steps obs obs"] $\Psi_0, \ldots, \Psi_H$ (or $\Theta_h$, or their running sums) with the step axis at ``-3``: the package's "time at ``-2``" convention counts from the trailing event dims, and the event here is an ``(obs, obs)`` matrix. Reshape to ``(*batch, steps, obs * obs)`` to plot the grid with `~~numpyro_forecast.convert.predictions_to_datatree()`. Raises ------ ValueError If ``horizon`` is negative. Notes ----- The posterior mean of the impulse responses is not the impulse response at the posterior mean of ``phi``: the recursion is nonlinear in ``phi``, and a single explosive draw dominates the mean at long horizons, so summarize draws with quantiles and check stability first. References ---------- Lütkepohl, H. (2005). *New Introduction to Multiple Time Series Analysis*, chapter 2. Springer. Orduz, J. *Bayesian VAR in NumPyro*, [juanitorduz.github.io/var_numpyro](https://juanitorduz.github.io/var_numpyro/), whose ``compute_irf`` this function generalizes. Pinder, T. *Impulso*, [github.com/thomaspinder/impulso](https://github.com/thomaspinder/impulso): the broadcast-over-draws recursion of its ``compute_ma_phi`` is the shape of this implementation, ported to a JAX scan. ## Priors Shrinkage prior moments for coefficient arrays. minnesota_prior(n_lags: int, n_obs: int, tightness: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, *, cross_shrinkage: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex = 0.5, decay: Literal['harmonic', 'geometric'] = 'harmonic', own_lag_mean: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex = 1.0) -> tuple[jaxtyping.Float[Array, 'lags obs obs'], jaxtyping.Float[Array, 'lags obs obs']] Return the Minnesota (Litterman) prior moments for VAR coefficients. The prior shrinks the coefficient of variable $j$ at lag $l$ in the equation of variable $i$ toward a mean that is nonzero only for the first own lag, with a standard deviation that tightens on longer lags and on cross-variable lags: $$ m_{l, ij} = \begin{cases} m_{\text{own}} & l = 1,\ i = j \\ 0 & \text{otherwise} \end{cases}, \qquad s_{l, ij} = \lambda \, d(l) \, \begin{cases} 1 & i = j \\ \kappa & i \neq j \end{cases}, $$ with $d(l) = 1/l$ (``decay="harmonic"``) or $d(l) = 1/l^2$ (``decay="geometric"``, the label Impulso uses; in Doan, Litterman and Sims it is the harmonic decay with exponent two). $m_{\text{own}} = 1$ is the random-walk belief for series in levels; use ``own_lag_mean=0.0`` for differenced or otherwise stationary series. The parameterization follows Impulso's ``MinnesotaPrior`` (fixed tightness rather than an estimated one, no residual-scale ratios), with ``own_lag_mean`` as the one addition. The classic Litterman scale ratio for the coefficient of variable $j$ in equation $i$ is $\sigma_i / \sigma_j$; apply it, when the series are on different scales, as ``scale * (sigma[:, None] / sigma[None, :])``. Parameters ---------- n_lags Number of lags $p \ge 1$. n_obs Number of series $k \ge 1$. tightness Overall shrinkage $\lambda > 0$ (the standard deviation of the first own lag). A modeling choice, so it has no default; a jax scalar is accepted, which lets a model sample it. cross_shrinkage Relative shrinkage $\kappa \in [0, 1]$ of cross-variable lags versus own lags (``1.0`` treats them alike, ``0.0`` pins them to the mean). decay Lag decay $d(l)$: ``"harmonic"`` for $1/l$, ``"geometric"`` for $1/l^2$. own_lag_mean Prior mean $m_{\text{own}}$ of the first own lag. Returns ------- loc : Float[Array, "lags obs obs"] Prior means in the ``phi`` layout of `~~numpyro_forecast.var.var_step()`. scale : Float[Array, "lags obs obs"] Prior standard deviations in the same layout. Raises ------ ValueError If ``n_lags`` or ``n_obs`` is below one, or a Python-number ``tightness`` is not positive or ``cross_shrinkage`` is outside ``[0, 1]`` (an unknown ``decay`` is a type error under the package's runtime type checking). Examples -------- ```python loc, scale = minnesota_prior(n_lags=2, n_obs=3, tightness=0.5, own_lag_mean=0.0) phi = numpyro.sample("phi", dist.Normal(loc, scale).to_event(3)) ``` References ---------- Litterman, R. B. (1986). Forecasting with Bayesian vector autoregressions: five years of experience. *Journal of Business & Economic Statistics*, 4(1). Doan, T., Litterman, R. B. and Sims, C. A. (1984). Forecasting and conditional projection using realistic prior distributions. *Econometric Reviews*, 3(1). Pinder, T. *Impulso*, ``MinnesotaPrior`` ([documentation](https://thomaspinder.github.io/Impulso/reference/generated/impulso.priors.MinnesotaPrior.html), [repository](https://github.com/thomaspinder/impulso)). ## Distribution surgery Time-axis operations on observation distributions, extensible via singledispatch. shift_loc(noise_dist: numpyro.distributions.distribution.Distribution, loc: jax.Array) -> numpyro.distributions.distribution.Distribution Re-center a zero-centered noise distribution at ``loc``. This converts Pyro's ``obs = data - prediction`` idiom into an additive shift of the observation distribution's location. Parameters ---------- noise_dist A zero-centered location-family distribution. loc The deterministic mean to add to the distribution's location. Returns ------- dist.Distribution A distribution centered at ``loc``. Raises ------ NotImplementedError If ``noise_dist`` is of an unsupported type. slice_time(noise_dist: numpyro.distributions.distribution.Distribution, index: slice) -> numpyro.distributions.distribution.Distribution Slice an elementwise distribution along the time axis ``-2``. The default implementation handles registered elementwise families (empty ``event_shape``, ``batch_shape`` ending in ``(time, obs)``; e.g. ``Normal``, ``StudentT``, ``Poisson``) by slicing each broadcast parameter. Correlated families register a dedicated dispatch instead. Parameters ---------- noise_dist The distribution to slice. index A ``slice`` applied to the time axis ``-2`` of the batch shape. Returns ------- dist.Distribution The same distribution family restricted to the selected time steps. Raises ------ NotImplementedError If ``noise_dist`` is not a registered elementwise family (and has no dedicated dispatch). prefix_condition(noise_dist: numpyro.distributions.distribution.Distribution, data: jax.Array) -> numpyro.distributions.distribution.Distribution Condition a ``(t+f)``-length distribution on a ``t``-length data prefix. For independent-over-time noise (the elementwise default) the conditional reduces to the forecast-horizon marginal, i.e. a time slice ``[t:]``. Only registered elementwise families take this path; correlated families (e.g. ``MultivariateNormal``) need a dedicated dispatch implementing a genuine Gaussian conditional, which is checked here rather than silently reduced to a slice (the R1 fix). Parameters ---------- noise_dist The observation distribution over the full horizon ``(*batch, t+f, obs)``. data The observed prefix with shape ``(*batch, t, obs)``. Returns ------- dist.Distribution The forecast-horizon distribution over ``(*batch, f, obs)``. Raises ------ NotImplementedError If ``noise_dist`` is not a registered elementwise family (and has no dedicated dispatch). register_elementwise(cls: type[numpyro.distributions.distribution.Distribution]) -> type[numpyro.distributions.distribution.Distribution] Declare a distribution family elementwise (usable as a decorator). An elementwise family is independent across every batch cell with an empty event shape, so `slice_time()` may slice its broadcast parameters along the time axis and `prefix_condition()` may reduce to the horizon marginal. Membership is by exact type; register each concrete subclass you rely on. Parameters ---------- cls The distribution class to register. Returns ------- type[dist.Distribution] ``cls`` unchanged, so this can decorate a class definition. ## Reparameterization Time-axis reparameterization of in-sample latents (Haar / DCT), after Pyro's `time_reparam`. time_reparam(model: numpyro_forecast.typing.ForecastModel, transform: Literal['haar', 'dct']) -> numpyro_forecast.typing.ForecastModel Reparameterize every in-sample time latent of ``model`` along the time axis. Port of the ``time_reparam`` argument of Pyro's ``Forecaster`` and ``HMCForecaster`` ([forecaster.py](https://github.com/pyro-ppl/pyro/blob/dev/pyro/contrib/forecast/forecaster.py)). The returned model is ``numpyro.handlers.reparam(model, config)`` with a config that targets every non-observed continuous sample site under the ``time`` plate opened by `~~numpyro_forecast.models.innovations()`, exactly as Pyro's ``time_reparam_haar`` targets every site inside its ``time`` plate. Each targeted site ``name`` becomes a ``deterministic`` site of the same shape, computed from a new sample site ``f"{name}_haar"`` or ``f"{name}_dct"`` whose event shape owns the time axis and every plate axis to its right (plate axes to the left of time stay batch axes). The transform is orthonormal, so the model's log density is unchanged; only the geometry inference sees changes. Parameters ---------- model A forecasting model ``(covariates, data=None) -> None``. transform ``"haar"`` for `numpyro.distributions.transforms.HaarTransform` (a multi-resolution average/difference basis, suited to blocky or multi-scale dynamics) or ``"dct"`` for `numpyro.distributions.transforms.DiscreteCosineTransform` (a cosine frequency basis, close to the Karhunen-Loeve basis of first-order Markov processes). Note that Pyro's own string mapping is swapped: its ``"haar"`` runs a DCT and its ``"dct"`` runs a Haar transform. Returns ------- ForecastModel The wrapped model. It is a `numpyro.handlers.reparam` handler and is the single object to hand to the guide, ``SVI`` / ``MCMC`` / blackjax, `~~numpyro_forecast.predictive.forecast()`, `~~numpyro_forecast.predictive.predict_in_sample()`, `~~numpyro_forecast.convert.to_datatree()` and the ``model_fn`` of `~~numpyro_forecast.evaluate.backtest()`. Raises ------ ValueError If ``model`` is already the result of `time_reparam()`. Nesting is not supported: the inner handler would turn the site into a deterministic before the outer one sees it, so the outer transform would silently be a no-op. Notes ----- - Create the wrapped model once and reuse it: the drivers jit-compile with the model as a static argument, so wrapping again inside a loop recompiles. - Apply it innermost. ``scope(time_reparam(model), prefix="a")`` yields ``a/drift_haar``; ``time_reparam(scope(model, "a"))`` yields ``a/a/drift_haar`` because the auxiliary sample passes through ``scope`` a second time (generic NumPyro reparam-under-scope behavior). - It composes with the per-block ``reparam=`` hook: after ``innovations(..., reparam=LocScaleReparam(0))`` the site under the plate is ``drift_decentered``, so the auxiliary site is ``drift_decentered_haar`` and both ``drift_decentered`` and ``drift`` become deterministic. This matches Pyro, whose config applies to every site in the plate. - Untouched sites: the ``_future`` suffix sites (they stay prior-drawn under ``time_future``, so `~~numpyro_forecast.predictive.forecast()` is unaffected), the scan sites of `~~numpyro_forecast.models.markov_series()`, the error sites of `~~numpyro_forecast.models.ssoe()`, observed sites and discrete sites. - Posterior dictionaries from `~~numpyro_forecast.predictive.draw_posterior()` and ``mcmc.get_samples()`` contain both ``drift`` (deterministic) and ``drift_haar``; ``Predictive`` substitutes only the latter and recomputes the former. ``init_to_value`` must therefore target ``drift_haar``. - Measured on random-walk level models: mean-field ``AutoNormal`` reaches a better ELBO for both transforms (DCT slightly ahead of Haar), but an optimizer schedule tuned for the original coordinates does not transfer as is. The rotated posterior is better conditioned, so it tolerates and may need a larger learning rate to converge within the same step budget; a schedule that is too small stalls with part of the intercept still held by the level. NUTS trajectories become cheaper (fewer leapfrog steps per iteration) while the effective sample size per draw is model dependent. Examples -------- ```python model_dct = time_reparam(seasonal_model, "dct") guide = AutoNormal(model_dct) svi = SVI(model_dct, guide, Adam(0.01), Trace_ELBO()) svi_result = svi.run(key_fit, 1_500, covariates[:t_obs], data, progress_bar=False) posterior = draw_posterior(key_post, guide, svi_result.params, num_samples=100) samples = forecast(key_pred, model_dct, posterior, data, covariates) ``` Literal(*args, **kwargs) ## Producing draws Drawing posterior samples and generating forecasts and in-sample predictions. draw_posterior(rng_key: jax.Array, guide: numpyro.infer.autoguide.AutoGuide, params: dict[str, jax.Array], num_samples: int, *, batch_size: int | None = None, device: jaxlib._jax.Device | str | None = None) -> dict[str, jax.Array | numpy.ndarray] Draw ``num_samples`` posterior samples of the latent sites from a fitted guide. The returned dict has the sample axis leading and is ready to pass to `forecast()` or NumPyro's ``Predictive``. An ``AutoDelta`` guide is a MAP point estimate: it is drawn once and tiled to ``num_samples`` (`_ensure_sample_axis_for_delta()`), since it carries no posterior spread of its own. Every other ``AutoGuide`` is sampled through a jitted, per-guide-cached ``sample_posterior`` (`_jitted_sample_posterior()`). Parameters ---------- rng_key PRNG key. guide The fitted variational guide, e.g. the ``AutoGuide`` instance passed to ``SVI``. params The learned variational parameters, e.g. the trained parameters from ``svi.run``'s result. num_samples Number of posterior draws. batch_size Optional chunk size for the drawing itself. Sampling a variational posterior materializes every latent and deterministic site for all draws at once, which on a wide panel is the largest allocation of the whole workflow. With ``batch_size`` set (strictly below ``num_samples``), the draws are sampled in chunks of exactly this many samples, each chunk is moved per ``device`` before the next is drawn, and the final chunk's overdraw is discarded, so accelerator memory is bounded by one chunk. Chunking changes the PRNG stream layout: draws are reproducible per ``(rng_key, batch_size)``. device Where each chunk of draws is moved as soon as it is drawn. ``"host"`` keeps every leaf in pageable host memory, so nothing of the result occupies accelerator memory: with the JAX CPU backend initialized it commits each leaf to ``jax.devices("cpu")[0]`` and returns committed `jax.Array` leaves (``np.asarray`` on one is a zero-copy view); without it (for example after ``numpyro.set_platform("cuda")``, or a ``JAX_PLATFORMS`` preset) it copies each chunk with `jax.device_get()` and returns NumPy arrays, since a CUDA client offers no pageable ``jax.Array`` container. It therefore needs no CPU backend and never pins memory. ``"numpy"`` forces the NumPy path; ``"pinned_host"`` commits to the accelerator's pinned host memory kind instead, a pool capped by ``XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB`` (64 GB by default on CUDA), so prefer ``"host"`` for large panels. A `jax.Device` or platform name like ``"cpu"`` commits the draws to that device (``"cpu"`` warns and takes the NumPy path when the CPU backend is missing). ``None`` keeps everything on the default device. ``device`` never changes the draw values. A host-committed jax result is not a drop-in replacement for a device array in your own ``jnp`` code: mixed with an uncommitted array an op runs on the CPU and returns a CPU-committed array, mixed with an accelerator-committed array it raises, and a pinned array raises on any mix. Feed a host posterior (in either container) straight into `forecast()`, `predict_in_sample()`, or `~~numpyro_forecast.convert.to_datatree()` (all of which already accept it), or convert explicitly first with ``np.asarray(x)`` (stays on host) or ``jax.device_put(x, device)`` (moves it to an accelerator). Returns ------- dict[str, Array | np.ndarray] Posterior samples of the latent sites, sample axis leading. With ``device="host"`` the leaves are committed to the CPU device, or NumPy arrays when no CPU backend is initialized. Raises ------ ValueError If ``num_samples`` or ``batch_size`` is not positive. HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._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 ----- For an MCMC fit, use its samples directly (``mcmc.get_samples()``); this function draws afresh from a variational guide, and chunks drawn from independent subkeys remain valid i.i.d. posterior samples. forecast(rng_key: jax.Array, model: numpyro_forecast.typing.ForecastModel, posterior: collections.abc.Mapping[str, jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex], data: jax.Array, covariates: jax.Array, *, batch_size: int | None = None, parallel: bool = True, device: jaxlib._jax.Device | str | None = None) -> jaxtyping.Num[Array, 'sample *batch future obs'] | jaxtyping.Num[ndarray, 'sample *batch future obs'] Sample forecasts for the steps in ``[t, duration)`` from a posterior. 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.predictive.draw_posterior()`). Parameters ---------- rng_key PRNG key. model The forecasting model callable (the same one that produced ``posterior``). posterior Posterior samples of the latent sites, sample axis leading. The output of a ``device="host"`` stage (CPU-committed jax leaves or NumPy leaves) is accepted directly. data Observed data with time at axis ``-2`` and length ``t``. covariates Covariates with time at axis ``-2`` and length ``duration > t``. batch_size Optional chunk size for sampling (caps peak memory). parallel 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 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, 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. The result feeds straight into `~~numpyro_forecast.convert.to_datatree()` and the ``batch_size``-chunked evaluation metrics in `~~numpyro_forecast.evaluate`, which accept host-resident draws. Returns ------- Num[Array, " 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.models.predict()`; with ``device="host"`` committed to the CPU device, or a NumPy array when no CPU backend is initialized). Raises ------ ValueError If ``covariates`` does not extend beyond ``data`` along the time axis. HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._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 ----- 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. predict_in_sample(rng_key: jax.Array, model: numpyro_forecast.typing.ForecastModel, posterior: collections.abc.Mapping[str, jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex], covariates: jax.Array, *, batch_size: int | None = None, parallel: bool = True, device: jaxlib._jax.Device | str | None = None) -> jaxtyping.Num[Array, 'sample *batch time obs'] | jaxtyping.Num[ndarray, 'sample *batch time obs'] Sample the in-sample posterior predictive of the ``obs`` site. 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 `~~numpyro_forecast.predictive.draw_posterior()`). Parameters ---------- rng_key PRNG key. model The forecasting model callable (the same one that produced ``posterior``). posterior Posterior samples of the latent sites, sample axis leading. The output of a ``device="host"`` stage (CPU-committed jax leaves or NumPy leaves) is accepted directly. covariates 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. batch_size Optional chunk size for sampling (caps peak memory). parallel 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 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 `~~numpyro_forecast.convert.to_datatree()` and the ``batch_size``-chunked evaluation metrics in `~~numpyro_forecast.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 ------ HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._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. ## Backtesting & evaluation Rolling-window backtesting and forecast metrics. backtest(rng_key: jax.Array, model_fn: collections.abc.Callable[[], numpyro_forecast.typing.ForecastModel], data: jax.Array, covariates: jax.Array, *, forecast_fn: numpyro_forecast.typing.ForecastFn, in_sample_fn: numpyro_forecast.typing.InSampleFn | None = None, metrics: collections.abc.Mapping[str, collections.abc.Callable[[jax.Array, jax.Array], jax.Array]] | None = None, per_window_metrics: collections.abc.Callable[[int, int, int], collections.abc.Mapping[str, collections.abc.Callable[[jax.Array, jax.Array], jax.Array]]] | None = None, transform: collections.abc.Callable[[jax.Array, jax.Array], tuple[jax.Array, jax.Array]] | None = None, window_type: Literal['expanding', 'rolling'] | None = None, train_window: int | None = None, min_train_window: int = 1, test_window: int | None = None, min_test_window: int = 1, stride: int = 1, num_samples: int = 100, batch_size: int | None = None, eval_train: bool = False, keep_predictions: bool = False, reuse_model: bool = True) -> list[numpyro_forecast.evaluate.BacktestResult] Backtest a forecasting model on a moving window of ``(train, test)`` data. Fitting and forecasting are delegated entirely to user-supplied closures rather than an OOP forecaster: ``forecast_fn`` fits ``model`` on the training window and forecasts the test horizon, and the optional ``in_sample_fn`` fits and scores the in-sample fit. Both closures own their own inference backend (SVI, MCMC, or anything else), so ``backtest`` itself has no dependency on how a model is fit. ``forecast_fn`` has the call signature (see `~~numpyro_forecast.typing.ForecastFn`): ```python forecast_fn( rng_key, model, train_data, train_covariates, full_covariates, num_samples, *, batch_size=None, ) -> draws # shape (num_samples, *batch, t2 - t1, obs) ``` where ``full_covariates`` spans the *full* window, ``covariates[..., t0:t2, :]`` (train followed by test), matching what the model needs to run the forecast horizon. The optional ``in_sample_fn`` has the call signature (see `~~numpyro_forecast.typing.InSampleFn`): ```python in_sample_fn( rng_key, model, train_data, train_covariates, num_samples, *, batch_size=None, ) -> draws # shape (num_samples, *batch, t1 - t0, obs) ``` ``batch_size`` is forwarded unchanged into both closures so a chunked implementation can bound its own device memory. A closure may return draws committed to host memory (e.g. via ``device="host"``, to cap peak accelerator usage): every metric in `DEFAULT_METRICS` accepts a host-committed ``pred`` or ``truth`` (or both), in any mix and regardless of ``batch_size``, moving a host-committed operand to device memory first where needed. Returning draws already on-device still avoids the extra host-to-device hop for metrics scored every window. A minimal ``forecast_fn`` built on plain NumPyro (``AutoNormal`` + ``SVI.run`` + ``Predictive``): ```python import numpyro from jax import random from numpyro.infer import SVI, Predictive, Trace_ELBO from numpyro.infer.autoguide import AutoNormal def forecast_fn( rng_key, model, train_data, train_covariates, full_covariates, num_samples, *, batch_size=None, ): guide = AutoNormal(model) svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO()) key_fit, key_post, key_pred = random.split(rng_key, 3) state = svi.run(key_fit, 1_000, train_covariates, train_data, progress_bar=False) posterior = guide.sample_posterior(key_post, state.params, sample_shape=(num_samples,)) predictive = Predictive(model, posterior_samples=posterior, return_sites=["forecast"]) return predictive(key_pred, full_covariates, train_data)["forecast"] ``` Parameters ---------- rng_key Base PRNG key (used for every window, matching Pyro). model_fn Factory returning a fresh `~~numpyro_forecast.typing.ForecastModel` per window. data Dataset with time at axis ``-2``. covariates Covariates with time at axis ``-2`` (same duration as ``data``). forecast_fn Closure that fits ``model`` on the training window and forecasts the test horizon (see `~~numpyro_forecast.typing.ForecastFn` and the contract above). in_sample_fn Optional closure that fits ``model`` on the training window and scores its in-sample fit (see `~~numpyro_forecast.typing.InSampleFn` and the contract above). Required when ``eval_train=True``. metrics Mapping of metric name to function; defaults to `DEFAULT_METRICS`. Each function takes ``(pred, truth)`` and returns a scalar array (see `~~numpyro_forecast.typing.Metric`); bind any metric-specific parameters with `functools.partial()`, e.g. ``{**DEFAULT_METRICS, "coverage": partial(eval_coverage, alpha=0.8)}``. per_window_metrics Optional ``(t0, t1, t2) -> Mapping[str, Metric]`` callable producing extra metrics merged over ``metrics`` for each window. Use it for window-dependent metrics such as a MASE scaled by that window's training data (`numpyro_forecast.metrics.make_mase()`). transform Optional ``(pred, truth) -> (pred, truth)`` applied before metrics. It runs before the metrics and receives the forecast/in-sample closure's draws as-is, so a transform that does its own array math against a device-resident operand must convert a host-committed ``pred`` or ``truth`` first, e.g. with `numpy.asarray()` (or move it back onto a device explicitly). window_type Windowing strategy. If ``None`` (default) it is inferred from ``train_window``: ``"expanding"`` when ``train_window`` is ``None`` and ``"rolling"`` when it is set, matching the historical behavior. Pass ``"expanding"`` to always train on all history from ``t0 = 0``, or ``"rolling"`` to hold the training length fixed at ``train_window`` and slide it forward. ``"expanding"`` and ``train_window`` are mutually exclusive, and ``"rolling"`` requires ``train_window`` (both validated). train_window Training window size; if ``None`` the window expands from the start. Required for ``window_type="rolling"``. min_train_window Minimum training window size for the expanding strategy (used when ``train_window`` is ``None``). test_window Test window size; if ``None`` forecasts to the end of the data. min_test_window Minimum test window size when ``test_window`` is ``None``. stride Step between successive train/test splits. num_samples Number of forecast samples per window. batch_size Optional chunk size forwarded to ``forecast_fn`` and ``in_sample_fn`` (see the contract above). eval_train If ``True``, also score the in-sample posterior predictive over each training window with the same ``metrics`` and store them in ``BacktestResult.train_metrics``. Requires ``in_sample_fn``. keep_predictions If ``True``, store each window's out-of-sample forecast samples (after ``transform``) on ``BacktestResult.prediction``. Defaults to ``False`` to avoid retaining large Monte Carlo arrays. reuse_model When ``True`` (default) and the windowing strategy is rolling, the model instance returned by the first ``model_fn()`` call is reused for every window so forecast/predict kernels can cache across windows. SVI still recompiles per window; for a single fused fit over all windows use `backtest_vectorized()`. Ignored for expanding windows and when ``False``. Returns ------- list[BacktestResult] One result per backtest window. Raises ------ ValueError If ``data`` and ``covariates`` durations differ, or if ``eval_train=True`` but ``in_sample_fn`` is ``None``. BacktestWindowError If ``window_type`` and ``train_window`` disagree (``"rolling"`` without a ``train_window``, or ``"expanding"`` with one). backtest_vectorized(rng_key: jax.Array, model_fn: collections.abc.Callable[[], numpyro_forecast.typing.ForecastModel], data: jax.Array, covariates: jax.Array, *, train_window: int, test_window: int, guide: numpyro_forecast.typing.Guide, optim: '_NumPyroOptim | None' = None, stride: int = 1, num_steps: int = 1001, num_samples: int = 100, metrics: collections.abc.Mapping[str, collections.abc.Callable[[jax.Array, jax.Array], jax.Array]] | None = None, keep_predictions: bool = False) -> numpyro_forecast.evaluate.VectorizedBacktestResult Rolling-window backtest with all windows fitted in one vmapped SVI run. Estimator-equivalent to `backtest()` with rolling windows; it differs only in PRNG stream layout and float reduction order, so the equivalence is statistical, not bitwise. Model, guide, and SVI compile once regardless of the number of windows, giving order-of-magnitude wall-clock wins for tens of windows on small models. PRNG: ``fold_in(rng_key, -1)`` seeds a discarded eager warm-up init; ``fold_in(rng_key, i)`` is the window-``i`` parent, split into SVI-init, posterior-draw, and forecast subkeys (the parent itself is never consumed; see `_window_key_streams()`). Parameters ---------- rng_key Base PRNG key. model_fn Factory returning a fresh model; called exactly once (per-window model variation is unsupported here, use `backtest()`). Must return the same model object ``guide`` was built on (see ``guide`` below). data Dataset with time at axis ``-2``. covariates Covariates with time at axis ``-2`` (same duration as ``data``). train_window Fixed training-window length (``>= 1``). test_window Fixed test-window length (``>= 1``). guide Any guide for the model object ``model_fn`` returns (see `~~numpyro_forecast.typing.Guide`): an autoguide instance built on that same object, sampled through its ``sample_posterior``, or a hand-written guide function with the model's signature, sampled through ``numpyro.infer.Predictive(guide, params=...)`` on each window's training data. Either way one guide is shared by every window: ```python model = make_model() backtest_vectorized(..., model_fn=lambda: model, guide=AutoNormal(model)) backtest_vectorized(..., model_fn=lambda: model, guide=my_guide_fn) ``` optim NumPyro optimizer; ``None`` uses ``numpyro.optim.Adam(0.01)``. For an optax optimizer, wrap it with ``numpyro.optim.optax_to_numpyro`` first. stride Step between successive windows (``>= 1``). num_steps Number of SVI steps per window. num_samples Number of forecast samples drawn per window. metrics Mapping of metric name to function; defaults to `DEFAULT_METRICS`. Each metric is vmapped over the window axis, so any pure-JAX mapping, including partial-bound variants such as ``{**DEFAULT_METRICS, "coverage": partial(eval_coverage, alpha=0.5)}``, is scored inside the single fused computation. Host-side metrics are unsupported here; use `backtest()`, or ``keep_predictions=True`` and score on the host. keep_predictions If ``True``, retain the stacked forecast samples on the result. Returns ------- VectorizedBacktestResult The stacked per-window losses, metrics, and window indices. Raises ------ ValueError If ``data`` and ``covariates`` durations differ. BacktestWindowError If ``train_window``, ``test_window``, or ``stride`` is ``< 1``, or there is no room for a single window. VectorizedMetricError If a metric forces a host conversion under ``vmap`` (it is not a pure JAX function). BacktestResult(t0: int, t1: int, t2: int, num_samples: int, walltime: float, metrics: dict[str, float], train_metrics: dict[str, float] = , prediction: jax.Array | numpy.ndarray | None = None) -> None Per-window result of a `backtest()` run. Attributes ---------- t0, t1, t2 Train-begin, train/test split, and test-end time indices. num_samples Number of forecast samples drawn. walltime Wall-clock seconds for the window's timed ``forecast_fn`` call. metrics Mapping of metric name to value for the window. train_metrics Mapping of metric name to in-sample value for the window. Empty unless ``backtest`` was called with ``eval_train=True``. prediction Out-of-sample forecast samples for the window (sample axis leading), or ``None`` unless ``backtest`` was called with ``keep_predictions=True``. VectorizedBacktestResult(t0: jax.Array, t1: jax.Array, t2: jax.Array, num_samples: int, losses: jax.Array, metrics: dict[str, jax.Array], predictions: jax.Array | None = None) -> None Result of a `backtest_vectorized()` run (all windows at once). Unlike `BacktestResult` this holds every window's values stacked along a leading window axis, because the windows are fitted, drawn, and scored in single vmapped passes rather than one call each. There are no per-window walltimes (a single fused computation covers all windows). Attributes ---------- t0, t1, t2 Train-begin, train/test split, and test-end time indices, each an integer array of shape ``(num_windows,)``. num_samples Number of forecast samples drawn per window. losses SVI loss history with shape ``(num_windows, num_steps)``. metrics Mapping of metric name to a ``(num_windows,)`` array of per-window values. predictions Stacked out-of-sample forecast samples with shape ``(num_windows, num_samples, *batch, test_window, obs)``, or ``None`` unless ``keep_predictions=True``. Literal(*args, **kwargs) evaluate.DEFAULT_METRICS dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2) evaluate_forecast(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float, *, metrics: collections.abc.Mapping[str, collections.abc.Callable[[jax.Array, jax.Array], jax.Array]] | None = None) -> dict[str, float] Evaluate forecast samples against ground truth for several metrics at once. A one-call convenience that applies each metric in ``metrics`` to the same forecast samples and ground truth, converting to host floats at the end. It is the one-shot counterpart to `backtest()` and is also used internally by `backtest()` to score each rolling window. Metric-specific parameters live with the metric in the ``metrics`` mapping, not on this function. To tune a metric, bind its keyword with `functools.partial()`; for example, to score coverage at the 80% level: ```python from functools import partial metrics = {**DEFAULT_METRICS, "coverage": partial(eval_coverage, alpha=0.8)} evaluate_forecast(pred, truth, metrics=metrics) ``` Parameters ---------- pred Forecast samples with the sample axis first, shape ``(sample, *batch)``. truth Ground-truth values with shape ``(*batch)``. metrics Mapping of metric name to function; when ``None`` defaults to `DEFAULT_METRICS` (``mae``, ``rmse``, ``crps`` and ``coverage``). Each function takes ``(pred, truth)`` and returns a scalar array (see `~~numpyro_forecast.typing.Metric`); bind any extra parameters with `functools.partial()` (see above). Returns ------- dict[str, float] Each metric name mapped to its value. results_to_dataframe(results: 'Sequence[BacktestResult] | VectorizedBacktestResult') -> 'pd.DataFrame' Flatten backtest results into a tidy one-row-per-window ``DataFrame``. Columns are prefix-namespaced so metric and in-sample-metric names never collide: window metrics become ``metric_`` and in-sample metrics ``train_metric_``, alongside the window indices ``t0``/``t1``/``t2``, ``num_samples``, and ``walltime``. Forecast samples are excluded. Windows may carry different metric sets (e.g. via ``backtest(per_window_metrics=...)``); the union of columns is used and missing entries are left as ``NaN``. A `VectorizedBacktestResult` from `backtest_vectorized()` is also accepted; it produces the same ``metric_`` columns for the same metric set, but has no ``train_metric_*`` or ``walltime`` columns (a vectorized run has no per-window walltimes), which are simply absent. Parameters ---------- results A sequence of `BacktestResult` from `backtest()`, or a single `VectorizedBacktestResult` from `backtest_vectorized()`. Returns ------- pandas.DataFrame One row per window with the namespaced columns described above. Raises ------ ImportError If ``pandas`` is not installed (``pip install numpyro_forecast[dataframes]``). eval_crps(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float, *, batch_size: int | None = None) -> jax.Array Empirical CRPS averaged over all data elements. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`). Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). batch_size Optional number of flattened data cells (the product of the batch shape, e.g. time times series) evaluated on the accelerator per pass; the sample axis is never chunked. With host-resident inputs (e.g. draws sampled with ``device="host"``) this bounds accelerator memory by ``sample * batch_size`` values plus the CRPS sort workspace instead of the full panel, since both ``pred`` and ``truth`` are staged as NumPy views before chunking, whatever their own memory kind. Chunking only changes the summation order of the final mean (results are equal to float tolerance, not bitwise); at or above the cell count the single-pass path runs instead, which moves any host-committed operand back to device memory first (`~~numpyro_forecast._offload._device_view()`) so either ``pred`` or ``truth`` (or both) may be host-committed regardless of ``batch_size``. ``None`` (default) evaluates in one pass. Returns ------- Array The mean empirical CRPS as a scalar array. eval_mae(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float) -> jax.Array Mean absolute error using the forecast sample median as point estimate. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`). ``pred`` and ``truth`` are moved to device memory first (`~~numpyro_forecast._offload._device_view()`), so either (or both) may be host-committed, e.g. draws sampled with ``device="host"``. Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). Returns ------- Array The mean absolute error as a scalar array. eval_rmse(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float) -> jax.Array Root mean squared error using the forecast sample mean as point estimate. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`). ``pred`` and ``truth`` are moved to device memory first (`~~numpyro_forecast._offload._device_view()`), so either (or both) may be host-committed, e.g. draws sampled with ``device="host"``. Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). Returns ------- Array The root mean squared error as a scalar array. eval_coverage(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float, *, alpha: float = 0.9, batch_size: int | None = None) -> jax.Array Empirical coverage of the central ``alpha`` prediction interval. The central ``alpha`` interval is bounded by the ``(1 - alpha) / 2`` and ``1 - (1 - alpha) / 2`` quantiles of the forecast samples; the metric is the fraction of ground-truth values that fall inside it. A well-calibrated forecast has coverage close to ``alpha``. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`); bind a non-default level with ``functools.partial(eval_coverage, alpha=...)``. Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). alpha Nominal interval level in ``(0, 1)``; defaults to ``0.9``. batch_size Optional number of flattened data cells evaluated on the accelerator per pass (see `eval_crps()`; the sample axis is never chunked). Coverage is a count of exact 0/1 indicators, so chunking recovers the identical count; only the precision of the final division differs from the single pass. Either path accepts a host-committed ``pred`` or ``truth`` (or both); the single pass moves them to device memory first (`~~numpyro_forecast._offload._device_view()`). ``None`` (default) evaluates in one pass. Returns ------- Array The fraction of ground truth inside the central ``alpha`` interval, as a scalar array. Raises ------ ValueError If ``alpha`` is not strictly inside ``(0, 1)``. crps_empirical(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float) -> jaxtyping.Float[Array, '*batch'] Compute the empirical Continuous Ranked Probability Score (CRPS). The CRPS generalises the mean absolute error to probabilistic forecasts and is computed elementwise as $$ \mathrm{CRPS}(F, y) = \mathbb{E}|X - y| - \tfrac{1}{2}\,\mathbb{E}|X - X'|, $$ where $X, X'$ are independent draws from the forecast distribution $F$. The expectations are estimated from the forecast ``sample`` axis using the sorted-sample $O(n \log n)$ identity. Parameters ---------- pred Forecast samples with the sample axis first, shape ``(sample, *batch)``. May be host-committed (e.g. draws sampled with ``device="host"``), regardless of whether ``truth`` is: either operand is moved to device memory first (`~~numpyro_forecast._offload._device_view()`). truth Ground-truth values with shape ``(*batch)`` (broadcastable to ``pred``). Returns ------- Float[Array, "*batch"] Elementwise CRPS, one value per ``batch`` location. References ---------- Tilmann Gneiting, Adrian E. Raftery (2007). "Strictly Proper Scoring Rules, Prediction, and Estimation". *Journal of the American Statistical Association*. eval_pinball(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float, *, quantile: float = 0.5) -> jax.Array Mean pinball (quantile) loss of the forecast ``quantile``. The pinball loss for the forecast $\hat q$ of quantile $\tau$ is $\max(\tau (y - \hat q), (\tau - 1)(y - \hat q))$, averaged over all data elements. At ``quantile=0.5`` it is half the mean absolute error. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`); ``quantile`` is static so each level specializes its own branch. ``pred`` and ``truth`` are moved to device memory first (`~~numpyro_forecast._offload._device_view()`), so either (or both) may be host-committed, e.g. draws sampled with ``device="host"``. Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). quantile Target quantile in ``(0, 1)``. Returns ------- Array The mean pinball loss as a scalar array. Raises ------ ValueError If ``quantile`` is not strictly inside ``(0, 1)``. eval_interval_score(pred: jaxtyping.Float[Array, 'sample *batch'] | jaxtyping.Float[ndarray, 'sample *batch'], truth: jaxtyping.Float[Array, '*batch'] | jaxtyping.Float[ndarray, '*batch'] | numpy.number | float, *, alpha: float = 0.9) -> jax.Array Mean Winkler interval score for the central ``alpha`` prediction interval. For the central ``alpha`` interval $[l, u]$ (the $(1-\alpha)/2$ and $1-(1-\alpha)/2$ quantiles), the interval score is $(u - l) + \tfrac{2}{1-\alpha}\big[(l - y)\mathbf 1_{yu}\big]$, averaged over all data elements. It rewards narrow intervals and penalizes ground truth falling outside them; lower is better. A pure JAX scalar kernel (see `~~numpyro_forecast.typing.Metric`). ``pred`` and ``truth`` are moved to device memory first (`~~numpyro_forecast._offload._device_view()`), so either (or both) may be host-committed, e.g. draws sampled with ``device="host"``. Parameters ---------- pred Forecast samples with the sample axis first. truth Ground-truth values (matching ``pred`` without the sample axis). alpha Nominal interval level in ``(0, 1)``. Returns ------- Array The mean interval score as a scalar array. Raises ------ ValueError If ``alpha`` is not strictly inside ``(0, 1)``. make_mase(train_data: jaxtyping.Float[Array, '*batch time obs_dim'] | jaxtyping.Float[ndarray, '*batch time obs_dim'], *, seasonality: int = 1) -> collections.abc.Callable[[jax.Array, jax.Array], jax.Array] Build a Mean Absolute Scaled Error metric scaled by ``train_data``. MASE divides the forecast MAE (using the sample median as point estimate) by the in-sample MAE of the seasonal-naive forecast on ``train_data``, ``mean(|y_t - y_{t-seasonality}|)``. The scale is computed once at factory time; the returned metric has the standard scalar-array signature (see `~~numpyro_forecast.typing.Metric`). ``train_data`` and the returned metric's ``pred``/``truth`` are all moved to device memory first (`~~numpyro_forecast._offload._device_view()`), so any of them may be host-committed, e.g. draws sampled with ``device="host"``. Parameters ---------- train_data Training data with time at axis ``-2``; leading batch axes are allowed and the seasonal-naive scale is averaged over all axes. seasonality Seasonal period (``>= 1``); ``1`` is the random-walk naive baseline. Returns ------- Metric A ``(pred, truth)`` callable computing MASE as a scalar array. Raises ------ ValueError If ``seasonality < 1``, ``train_data`` is not longer than ``seasonality`` along the time axis, or the seasonal-naive scale is zero (a constant series, for which MASE is undefined). ## ArviZ export Convert posteriors into ArviZ-schema xarray DataTrees for diagnostics and plotting. to_datatree(rng_key: jax.Array, model: numpyro_forecast.typing.ForecastModel, posterior: collections.abc.Mapping[str, jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex], data: jax.Array, covariates: jax.Array, *, num_chains: int = 1, batch_size: int | None = None, device: jaxlib._jax.Device | str | None = 'host', coords: collections.abc.Mapping[str, collections.abc.Sequence[Any]] | None = None, time_coord: collections.abc.Sequence[Any] | None = None, posterior_dims: collections.abc.Mapping[str, collections.abc.Sequence[str]] | None = None, covariate_dims: collections.abc.Sequence[str] | None = None) -> 'xarray.DataTree' Convert an already-drawn posterior into an ArviZ-schema `xarray.DataTree`. Posterior-first: callers draw their own posterior (``mcmc.get_samples()`` for MCMC, `~~numpyro_forecast.predictive.draw_posterior()` for a variational fit) and pass it in; ``to_datatree`` never draws a posterior of its own. ``rng_key`` is consumed only by the in-sample posterior-predictive draws and, when a forecast horizon is present, the forecast draws. Parameters ---------- rng_key PRNG key for the in-sample predictive draws and, when a horizon is present, the forecast draws. model The forecasting model that produced ``posterior``. posterior Posterior samples of the latent sites, with a single flattened sample axis leading (NumPyro's ``mcmc.get_samples()`` order, or the output of `~~numpyro_forecast.predictive.draw_posterior()`). Host-committed leaves (the output of ``draw_posterior(..., device="host")``) and NumPy leaves are accepted directly. data In-sample data with time at axis ``-2``. covariates Covariates with time at axis ``-2``. When ``covariates`` extends beyond ``data`` along the time axis (the package-wide shape convention for a forecast horizon), the trailing rows are treated as future covariates: the returned tree additionally carries ``predictions`` (forecast ``obs`` draws from `~~numpyro_forecast.predictive.forecast()`) and ``predictions_constant_data`` groups. num_chains Number of chains to split ``posterior``'s flattened sample axis into (and, identically, the in-sample/forecast predictive draws, which are drawn with the same sample count). Defaults to ``1`` (a single pseudo-chain, correct for a posterior with no chain structure, e.g. SVI or Pathfinder draws). For an MCMC posterior, pass the ``num_chains`` the sampler was run with; see `_reshape_chains()` for the reshape contract and its divisibility requirement. batch_size Optional chunk size that bounds how many draws touch the accelerator at once, across both the in-sample and forecast predictive sampling. When set, sampling runs in chunks of this many draws, each chunk moved to ``device`` before the next is drawn. The per-chunk accelerator footprint is a handful of ``(batch_size, time, series)`` buffers, so it scales linearly with this value times the panel width: on wide panels lower it until a chunk fits. The batch size must be strictly below the draw count for that bound to hold: at or above it, sampling falls back to the single-shot path and the full array is materialized on the default device before the single transfer. Chunking changes the PRNG stream layout of the predictive draws, so results are reproducible per ``(rng_key, batch_size)``. ``None`` (default) samples everything in one shot (the results are still moved to ``device``). device Where the predictive draws are moved as they are sampled, forwarded to the ``device`` argument of `~~numpyro_forecast.predictive.predict_in_sample()` and `~~numpyro_forecast.predictive.forecast()` (the placement contract of `~~numpyro_forecast.predictive.draw_posterior()`). It is resolved once and the same placement is handed to both, so an unmet ``"cpu"`` warns once per export. The default ``"host"`` keeps the predictive draws in pageable host memory (jax Arrays the tree views as NumPy without a copy, or NumPy arrays when no CPU backend is initialized), which is what bounds accelerator memory when ``batch_size`` is set; pass ``None`` to keep the draws on the default device (chunked compute without per-chunk host transfers, for when the draws fit on the accelerator and transfers would dominate runtime). coords Optional extra coordinates; these take precedence over the generated ``time`` coordinate. They also propagate to the forecast groups, where the generated forecast ``time`` takes precedence instead (a user ``time`` entry covers the in-sample window; use ``time_coord`` for explicit forecast time values). time_coord Optional explicit time coordinate values. Without a forecast horizon it covers the in-sample window (defaults to ``range(n_time)``); with a horizon it must cover the full ``covariates`` length and is split into the in-sample and forecast time coordinates (the default is the integer continuation). posterior_dims Optional mapping from a posterior site name to its non-sample dimension names, e.g. ``{"drift": ["time"]}``. Sites listed here share the tree-wide ``time`` coordinate; unlisted sites keep ArviZ's auto-named dims. This is an explicit opt-in on purpose: inferring time-indexed sites from trace shapes is fragile (a coincidental ``n_params == n_time`` would misattribute the axis). covariate_dims Optional dimension names for the stored covariates, one per axis; defaults to the 2-D ``("time", "covariate_dim")`` layout. Use this when ``covariates`` carries extra batch axes, e.g. a panel tensor shaped ``(channel, time, series)`` with ``covariate_dims=["channel", "time", "series"]``. The time axis is always ``-2`` (the package-wide convention), so its entry should be named ``"time"`` to share the tree-wide time coordinate. Returns ------- xarray.DataTree A tree with ``posterior`` (``(chain, draw, ...)``, split per ``num_chains``), ``posterior_predictive`` (in-sample ``obs``), ``observed_data``, and ``constant_data`` groups. When ``covariates`` extends beyond ``data``, also ``predictions`` and ``predictions_constant_data`` groups (sharing the same ``num_chains`` split). Raises ------ ValueError If ``covariates`` is shorter than ``data`` along the time axis, if ``time_coord`` is given but its length does not match the in-sample window plus the forecast horizon, or if ``posterior``'s sample count is not evenly divisible by ``num_chains``. CovariateDimsError If ``covariate_dims`` does not name every ``covariates`` axis. HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._resolve_device()`). Warns ----- UserWarning If ``device="cpu"`` is requested and the JAX CPU backend is not initialized, so the predictive draws take the NumPy path of ``"host"`` instead (once per call). Notes ----- When a forecast horizon is present, ``rng_key`` is split internally into a predictive subkey and a forecast subkey, so passing the same key twice never correlates the two sample sets. When there is no horizon, ``rng_key`` is used unsplit for the in-sample predictive draw. ``batch_size`` is the built-in route to memory-bounded predictive sampling; for fully manual control over the forecast draws, build the in-sample tree with matching-length covariates and attach the horizon with `add_forecast_groups()`. add_forecast_groups(tree: 'xarray.DataTree', forecast_samples: jaxtyping.Num[Array, 'sample future obs'] | jaxtyping.Num[ndarray, 'sample future obs'], covariates_future: jax.Array, *, time_coord: collections.abc.Sequence[Any] | None = None, covariate_dims: collections.abc.Sequence[str] | None = None) -> 'xarray.DataTree' Attach out-of-sample forecast groups to a copy of ``tree``. Adds a ``predictions`` group (the forecast ``obs`` draws) and a ``predictions_constant_data`` group (the future covariates). The forecast ``time`` coordinate continues the in-sample one: integer continuation by default, or explicit values via ``time_coord``. This is the step-by-step route for draws you produced yourself; `to_datatree()` attaches the same groups automatically when its ``covariates`` extend beyond ``data``. Parameters ---------- tree A tree from `to_datatree()` (its ``observed_data`` time coordinate is continued). forecast_samples Forecast draws shaped ``(num_samples, future, obs)`` from `~~numpyro_forecast.predictive.forecast()`. covariates_future Future covariates shaped ``(future, covariate_dim)``, or any layout with time at axis ``-2`` when ``covariate_dims`` names the axes. time_coord Optional explicit forecast time coordinate; defaults to integer continuation of the in-sample time. Required when the in-sample time coordinate is non-integer (e.g. datetime64): auto-continuing would have to guess the frequency, so explicit values are demanded instead. covariate_dims Optional dimension names for ``covariates_future``, one per axis. When omitted, the names are inherited from the tree's ``constant_data["covariates"]`` variable (falling back to ``("time", "covariate_dim")`` if the tree carries no stored covariates), so the forecast covariates always share the in-sample axis names. When given explicitly, the names must match the stored ones. See `to_datatree()`. Returns ------- xarray.DataTree A new tree with the ``predictions`` and ``predictions_constant_data`` groups added. Raises ------ ValueError If ``time_coord`` is given but its length differs from the forecast horizon, or if it is omitted while the in-sample time coordinate is non-integer. CovariateDimsError If the resolved ``covariate_dims`` (explicit or inherited) do not name every ``covariates_future`` axis, or if explicit names disagree with the dimension names already stored on the tree's ``constant_data["covariates"]``. predictions_to_datatree(predictions: jaxtyping.Float[Array, 'sample time series'] | jaxtyping.Float[ndarray, 'sample time series'], x: jaxtyping.Num[Array, 'time'] | jaxtyping.Num[ndarray, 'time'], series: collections.abc.Sequence[typing.Any], *, group: str = 'posterior_predictive', observed: jaxtyping.Float[Array, 'time series'] | jaxtyping.Float[ndarray, 'time series'] | None = None) -> 'xarray.DataTree' Pack prediction draws into a DataTree laid out for per-series ``plot_lm`` faceting. The array-level counterpart of `to_datatree()`: instead of a fit, it takes prediction draws from **any** predictive group (prior predictive, posterior predictive, or forecasts), possibly already transformed (rescaled to original units, clipped at zero, subset to a few series). The draws get a single pseudo-chain, and ``constant_data`` carries the independent variable ``"t"`` broadcast to ``(time, series)`` so that ``arviz.plot_lm(tree, y="obs", x="t", plot_dim="time", ...)`` facets one panel per series; band artists are then reachable via ``pc.viz["ci_band"]["t"]`` and axes via ``pc.get_target("t", {"series": label})``. ``plot_lm`` requires an ``observed_data`` group even when the observation scatter is disabled, so when ``observed`` is ``None`` a zeros placeholder is stored; it is never drawn under ``visuals={"observed_scatter": False}``. Parameters ---------- predictions Prediction draws with the sample axis first, shape ``(sample, time, series)``. x Independent-variable values, shape ``(time,)``. Must be numeric: ``plot_lm`` cannot draw ``datetime64`` values (it concatenates ``x`` with the float predictions internally), so pass `matplotlib.dates.date2num()` floats and re-format the tick labels with `matplotlib.dates.ConciseDateFormatter`. series One label per series, defining the ``series`` coordinate. group Predictive group to store the draws under (e.g. ``"prior_predictive"``, ``"posterior_predictive"``, ``"predictions"``). observed Optional observations, shape ``(time, series)``, stored in ``observed_data``; when ``None`` a zeros placeholder is stored instead. Returns ------- xarray.DataTree A tree with the ``group``, ``observed_data``, and ``constant_data`` groups; ``obs`` has dims ``(chain, draw, time, series)`` and ``t`` has dims ``(time, series)``. Raises ------ ValueError If ``series`` does not have one label per series in ``predictions``. ## Extensions (contrib) Optional backends behind pyproject extras (never imported by default). BlackjaxNUTSKernel(model: numpyro_forecast.typing.ForecastModel | None = None, *, num_adaptation_steps: int = 500, target_acceptance_rate: float = 0.8) -> None BlackJAX NUTS with Stan-style window adaptation. Adaptation (step size and inverse mass matrix) runs once in `_BlackjaxKernel.init()` via ``blackjax.window_adaptation``; each MCMC step is then a single tuned NUTS step. Run configuration ----------------- Pass this kernel to `numpyro.infer.MCMC` with **``chain_method="sequential"`` only**. With more than one chain, ``chain_method="vectorized"`` hands `_BlackjaxKernel.init()` a stacked ``rng_key`` of shape ``(num_chains, 2)`` rather than vmapping a per-chain call, and the ``jax.random.split`` inside ``init`` rejects a non-scalar key with a ``ValueError`` (pinned by the test suite). ``chain_method="parallel"`` is likewise unsupported: it is not exercised by this package and shares the same single-key assumption. Also pass **``num_warmup=0``**: adaptation runs once inside `_BlackjaxKernel.init()` (via ``blackjax.window_adaptation``), so any NumPyro-driven warmup steps on top of that are simply discarded work, not additional tuning. Parameters ---------- model The NumPyro model to sample from. num_adaptation_steps Number of window-adaptation steps. target_acceptance_rate Target acceptance probability for dual-averaging step-size adaptation. BlackjaxMCLMCKernel(model: numpyro_forecast.typing.ForecastModel | None = None, *, num_tuning_steps: int = 500) -> None BlackJAX Microcanonical Langevin Monte Carlo (MCLMC). The step size, trajectory length ``L``, and diagonal preconditioner (inverse mass matrix) are tuned once in `_BlackjaxKernel.init()` via ``blackjax.mclmc_find_L_and_step_size``; each MCMC step is then a single tuned MCLMC step. Run configuration ----------------- Pass this kernel to `numpyro.infer.MCMC` with **``chain_method="sequential"`` only**. With more than one chain, ``chain_method="vectorized"`` hands `_BlackjaxKernel.init()` a stacked ``rng_key`` of shape ``(num_chains, 2)`` rather than vmapping a per-chain call, and the ``jax.random.split`` inside ``init`` rejects a non-scalar key with a ``ValueError`` (pinned by the test suite). ``chain_method="parallel"`` is likewise unsupported: it is not exercised by this package and shares the same single-key assumption. Also pass **``num_warmup=0``**: tuning runs once inside `_BlackjaxKernel.init()` (via ``blackjax.mclmc_find_L_and_step_size``), so any NumPyro-driven warmup steps on top of that are simply discarded work, not additional tuning. Parameters ---------- model The NumPyro model to sample from. num_tuning_steps Number of tuning steps for ``L`` and the step size. BlackjaxCustomKernel(model: numpyro_forecast.typing.ForecastModel | None = None, *, build_fn: collections.abc.Callable[..., object]) -> None Adapt an arbitrary BlackJAX sampler via a user-supplied ``build_fn``. The escape hatch for kernels without a dedicated wrapper. ``build_fn`` receives ``(rng_key, logdensity_fn, position, num_warmup)`` and must return ``(inner_state, step_fn)``, where ``inner_state`` exposes ``.position`` over the same sites as the model and ``step_fn`` has signature ``(rng_key, inner_state) -> (inner_state, info)``. The base class validates the returned state's key set and raises `TypeError` if it is malformed. Run configuration ----------------- Pass this kernel to `numpyro.infer.MCMC` with **``chain_method="sequential"`` only**. With more than one chain, ``chain_method="vectorized"`` hands `_BlackjaxKernel.init()` a stacked ``rng_key`` of shape ``(num_chains, 2)`` rather than vmapping a per-chain call, and the ``jax.random.split`` inside ``init`` rejects a non-scalar key with a ``ValueError`` (pinned by the test suite). ``chain_method="parallel"`` is likewise unsupported: it is not exercised by this package and shares the same single-key assumption. Also pass **``num_warmup=0``**: whatever adaptation ``build_fn`` performs runs once inside `_BlackjaxKernel.init()`, so any NumPyro-driven warmup steps on top of that are simply discarded work, not additional tuning. Parameters ---------- model The NumPyro model to sample from. build_fn The blackjax build callable described above. PathfinderFit(state: Any, model: numpyro_forecast.typing.ForecastModel, covariates: jax.Array, data: jax.Array, elbo: float) -> None The result of fitting a forecasting model with BlackJAX Pathfinder. A plain-data (picklable) container: it holds the raw blackjax ``PathfinderState``, the model and its data/covariates (needed to rebuild the unconstrained-to-constrained transform when drawing), and the fitted ELBO. Draws are produced lazily by `pathfinder_samples()`. Attributes ---------- state The blackjax ``PathfinderState`` (the variational approximation). model The forecasting model that was fit. covariates In-sample covariates used at fit time (time at axis ``-2``). data In-sample data used at fit time (time at axis ``-2``). elbo The evidence lower bound of the fitted approximation. fit_pathfinder(rng_key: jax.Array, model: numpyro_forecast.typing.ForecastModel, data: jax.Array, covariates: jax.Array, *, num_elbo_samples: int = 200, ftol: float = 1e-05, maxiter: int = 30, maxcor: int = 10, maxls: int = 1000, gtol: float = 1e-08, init_params: dict[str, jax.Array] | None = None) -> numpyro_forecast.contrib.blackjax.PathfinderFit Fit a forecasting model with BlackJAX Pathfinder variational inference. PRNG: ``rng_key`` is split into a model-initialization stream and a Pathfinder-approximation stream. Parameters ---------- rng_key PRNG key for initialization and the Pathfinder run. model The forecasting model function ``(covariates, data=None)``. data In-sample data with time at axis ``-2``. covariates Covariates with time at axis ``-2`` and the same duration as ``data``. num_elbo_samples Number of Monte Carlo samples used to estimate the ELBO along the L-BFGS optimization path. ftol L-BFGS relative function-value tolerance (convergence criterion). maxiter Maximum number of L-BFGS iterations (blackjax's default is ``30``). Models with per-step time latents typically need far more: the default suits tens of parameters, while a few hundred parameters (e.g. a 400-step random walk) converge around ``1_000``. maxcor L-BFGS history size; caps the rank of the low-rank-plus-diagonal covariance correction at roughly ``2 * maxcor``. High-dimensional posteriors need it raised well above the default of ``10``. maxls Maximum number of line-search steps per L-BFGS iteration. gtol L-BFGS gradient-norm convergence tolerance. init_params Optional initial unconstrained position, overriding NumPyro's own initialization draw (``param_info.z``); ``initialize_model`` still runs regardless, since it also builds the potential/postprocess functions. Returns ------- PathfinderFit The fitted variational approximation. See Also -------- fit_multipathfinder The recommended entry point: runs several L-BFGS paths in parallel and combines them (see `multipathfinder_samples()`), instead of relying on this single path's local normal approximation. Notes ----- Runs with `_stable_bfgs_sample()` patched into blackjax (via `_ensure_stable_bfgs_sample()`): upstream's ``bfgs_sample`` underflows its log-determinant beyond a few hundred parameters, which floors every path ELBO at ``-inf`` and makes the returned state garbage. pathfinder_samples(rng_key: jax.Array, fit: numpyro_forecast.contrib.blackjax.PathfinderFit, num_samples: int, *, batch_size: int | None = None, device: jaxlib._jax.Device | str | None = None) -> dict[str, jax.Array | numpy.ndarray] Draw ``num_samples`` posterior samples from a fitted Pathfinder approximation. The returned dict has the sample axis leading and is ready to pass to `~~numpyro_forecast.predictive.forecast()` or NumPyro's ``Predictive``, exactly like `~~numpyro_forecast.predictive.draw_posterior()`. Chunking and device offload are delegated to the shared `~~numpyro_forecast._offload._draw_chunked()` driver, so the same memory-bounding contract applies. PRNG: within each chunk (the whole draw, when unchunked), the chunk key is split into a model-initialization stream and a Pathfinder-sampling stream (the init param draws are unused, but the streams are kept isolated). Parameters ---------- rng_key PRNG key. fit A fit from `fit_pathfinder()`. num_samples Number of posterior draws. batch_size Optional chunk size for the drawing itself; see `~~numpyro_forecast.predictive.draw_posterior()` (the same memory/reproducibility contract applies). device Where each chunk of draws is moved as soon as it is drawn; see `~~numpyro_forecast.predictive.draw_posterior()`, including for the NumPy path taken when the JAX CPU backend is not initialized and for the rules on mixing a host-committed result with other arrays. Returns ------- dict[str, Array] Posterior samples of the latent sites, sample axis leading (with ``device="host"``: leaves committed to the CPU device, or NumPy arrays when no CPU backend is initialized). Raises ------ ValueError If ``num_samples`` or ``batch_size`` is not positive. HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._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. MultiPathfinderFit(state: Any, model: numpyro_forecast.typing.ForecastModel, covariates: jax.Array, data: jax.Array, elbos: tuple[float, ...], log_weights: jax.Array, pareto_k: float) -> None The result of fitting a forecasting model with multi-path BlackJAX Pathfinder. A plain-data (picklable) container, mirroring `PathfinderFit`: it holds the raw blackjax ``MultipathfinderState`` (the pooled per-path approximations), the model and its data/covariates, the per-path ELBOs, and the Pareto-smoothed importance sampling (PSIS) weights/diagnostic over the pooled draws. Draws are produced lazily by `multipathfinder_samples()`. Attributes ---------- state The blackjax ``MultipathfinderState`` (per-path ``PathfinderState`` objects, the pooled samples, and their log target/approximation densities). model The forecasting model that was fit. covariates In-sample covariates used at fit time (time at axis ``-2``). data In-sample data used at fit time (time at axis ``-2``). elbos The evidence lower bound of each path's fitted approximation, in path order; converted eagerly with ``float()`` so pickling this fit never carries JAX tracers. log_weights Normalized PSIS log importance weights over the flattened pool of ``num_paths * num_elbo_samples`` draws (path-major, matching ``state.samples``); a fit-time diagnostic, valid only for that exact stored pool. `multipathfinder_samples()` draws fresh samples and recomputes its own weights, so it never consumes these. pareto_k The Pareto shape-parameter diagnostic for the fit-time PSIS weights: below ``0.5`` is reliable, ``0.5`` to ``0.7`` is borderline, and above ``0.7`` indicates the importance weights are unreliable. `multipathfinder_samples()` reads it to decide whether ``resample="auto"`` uses PSIS or ELBO-weighted path sampling. fit_multipathfinder(rng_key: jax.Array, model: numpyro_forecast.typing.ForecastModel, data: jax.Array, covariates: jax.Array, *, num_paths: int = 4, num_elbo_samples: int = 200, maxiter: int = 30, maxcor: int = 10, maxls: int = 1000, gtol: float = 1e-08, ftol: float = 1e-05, init_params: dict[str, jax.Array] | None = None) -> numpyro_forecast.contrib.blackjax.MultiPathfinderFit Fit a forecasting model with multi-path BlackJAX Pathfinder and PSIS resampling. Runs ``num_paths`` independent single-path Pathfinder approximations in parallel (vmapped L-BFGS runs, by default each from its own fresh ``init_to_uniform`` starting point) and scores the pooled per-path draws with Pareto-smoothed importance sampling (PSIS), so the returned fit is not tied to a single mode. This is the recommended entry point over `fit_pathfinder()`: a single L-BFGS path can settle in one mode and its local normal approximation may not reflect the rest of a multimodal or otherwise hard posterior. The PSIS weights computed here are a fit-time diagnostic; the draws themselves come from `multipathfinder_samples()`, which resamples fresh per-path draws and reads ``pareto_k`` to pick between PSIS and ELBO-weighted path sampling. PRNG: ``rng_key`` is split into a model-initialization stream and a multipath-approximation stream; the initialization stream is further split into one subkey per path so that (when ``init_params`` is not supplied) every path starts from its own independent ``init_to_uniform`` draw, exactly the diverse starting points multipath Pathfinder wants. Parameters ---------- rng_key PRNG key for initialization and the multipath Pathfinder run. model The forecasting model function ``(covariates, data=None)``. data In-sample data with time at axis ``-2``. covariates Covariates with time at axis ``-2`` and the same duration as ``data``. num_paths Number of independent L-BFGS paths to run in parallel (vmapped). num_elbo_samples Number of Monte Carlo samples drawn per path to estimate its ELBO and to build the pooled sample used for PSIS resampling. maxiter Maximum number of L-BFGS iterations per path. maxcor L-BFGS history size; caps the rank of the low-rank-plus-diagonal covariance correction at roughly ``2 * maxcor``. High-dimensional posteriors need it raised well above the default of ``10``. maxls Maximum number of line-search steps per L-BFGS iteration. gtol L-BFGS gradient-norm convergence tolerance. ftol L-BFGS relative function-value convergence tolerance. init_params Optional starting positions, one per path, overriding the default per-path ``init_to_uniform`` draws. Unlike the single-path `fit_pathfinder()` argument of the same name, every leaf must already carry a leading axis of size ``num_paths`` (validated; a mismatch raises ``ValueError``). Returns ------- MultiPathfinderFit The fitted multipath approximation together with its per-path ELBOs and the fit-time PSIS log weights/``pareto_k`` diagnostic over the pooled draws. Raises ------ ValueError If ``num_paths`` is not positive, or ``init_params`` is supplied without a leading axis of size ``num_paths`` on every leaf. Warns ----- UserWarning If one or more per-path ELBOs are non-finite (raise ``maxiter``/``maxcor`` or inspect ``fit.elbos``), or if the PSIS ``pareto_k`` diagnostic exceeds ``0.7``, in which case ``multipathfinder_samples(..., resample="auto")`` falls back to ELBO-weighted path sampling (increase ``num_paths``/``maxiter``/``maxcor``, or fall back to MCMC). Neither condition raises: the fit is always returned, carrying everything needed to inspect and diagnose it. Notes ----- Runs with `_stable_bfgs_sample()` patched into blackjax (via `_ensure_stable_bfgs_sample()`, called before ``blackjax.vi.multipathfinder.multi_approximate``): ``multi_approximate`` calls ``approximate``/``sample`` imported from ``blackjax.vi.pathfinder``, whose module globals hold ``bfgs_sample``, so patching that module (as `fit_pathfinder()` already does) covers the multipath route too. multipathfinder_samples(rng_key: jax.Array, fit: numpyro_forecast.contrib.blackjax.MultiPathfinderFit, num_samples: int, *, resample: Literal['auto', 'psis', 'elbo'] = 'auto', batch_size: int | None = None, device: jaxlib._jax.Device | str | None = None) -> dict[str, jax.Array | numpy.ndarray] Draw ``num_samples`` posterior samples from a fitted multipath Pathfinder fit. Every call draws fresh unconstrained samples from each path's fitted normal approximation (``num_samples`` per path, via ``blackjax.vi.pathfinder.sample`` vmapped over the paths) and then combines the paths into a single set of ``num_samples`` draws. Nothing is recycled from the small pool stored at fit time, so asking for more draws than that pool held costs nothing in duplication. The output contract is identical to `~~numpyro_forecast.predictive.draw_posterior()` and `pathfinder_samples()`: `~~numpyro_forecast.predictive.forecast()` and NumPyro's ``Predictive`` consume it unchanged. Chunking and device offload are delegated to the shared `~~numpyro_forecast._offload._draw_chunked()` driver, so the same memory-bounding contract applies. Resampling strategies --------------------- ``resample="psis"`` Pool the ``num_paths * num_samples`` fresh draws, score them under the model (``logp``) and under their own path's approximation (``logq``), Pareto-smooth the importance ratios, and resample ``num_samples`` of them with replacement. This is the textbook multipath Pathfinder estimator and the most faithful one *when the importance weights behave*. ``resample="elbo"`` Pick a whole path per draw, with probability ``softmax(fit.elbos)`` over the paths whose ELBO is finite (non-finite paths get zero weight, and an all-non-finite fit degrades to a uniform choice), then take one fresh draw from that path. With ELBO gaps of hundreds of nats this collapses onto the single best path, which is the right answer when the paths disagree strongly: no single draw is ever duplicated. ``resample="auto"`` (default) ``"psis"`` when ``fit.pareto_k <= 0.7``, ``"elbo"`` otherwise. The gate matters because importance sampling degenerates in high dimensions: on a posterior with hundreds of parameters the ratio ``logp - logq`` is dominated by a handful of draws, ``pareto_k`` climbs far above ``0.7``, and PSIS resampling collapses the answer onto those few draws. ELBO-weighted path sampling never concentrates like that, because it reweights *paths* (a handful of well-separated numbers) rather than individual draws. PRNG: within each chunk (the whole draw, when unchunked), the chunk key is split into a model-initialization stream (used only to rebuild the constraining transform), a per-path sampling stream, and a resampling-index stream. The split is the same in every mode, so ``resample="auto"`` produces bitwise the same draws as the explicit mode it resolves to. Parameters ---------- rng_key PRNG key. fit A fit from `fit_multipathfinder()`. num_samples Number of posterior draws. resample How to combine the per-path draws: ``"auto"`` (default), ``"psis"``, or ``"elbo"``; see "Resampling strategies" above. batch_size Optional chunk size for the drawing itself; see `~~numpyro_forecast.predictive.draw_posterior()` (the same memory/reproducibility contract applies). Note that each chunk draws ``num_paths * batch_size`` samples internally. device Where each chunk of draws is moved as soon as it is drawn; see `~~numpyro_forecast.predictive.draw_posterior()`, including for the NumPy path taken when the JAX CPU backend is not initialized and for the rules on mixing a host-committed result with other arrays. Returns ------- dict[str, Array] Posterior samples of the latent sites, sample axis leading (with ``device="host"``: leaves committed to the CPU device, or NumPy arrays when no CPU backend is initialized). Raises ------ ValueError If ``num_samples`` or ``batch_size`` is not positive, or ``resample`` is not one of ``"auto"``, ``"psis"``, ``"elbo"``. HostMemoryKindError If ``device="pinned_host"`` is requested on a device that exposes no host memory kind (see `~~numpyro_forecast._offload._host_memory_kind()`). DevicePlatformError If ``device`` names a platform whose backend is not initialized (see `~~numpyro_forecast._offload._resolve_device()`). Warns ----- UserWarning In PSIS mode, if the ``pareto_k`` recomputed on this call's fresh draws exceeds ``0.7``. This is a separate diagnostic from ``fit.pareto_k``, which is computed once over the pool stored at fit time. Also if ``device="cpu"`` is requested and the JAX CPU backend is not initialized, so the draws take the NumPy path of ``"host"`` instead. ## Typing Public type contracts. ForecastModel(*args, **kwargs) A NumPyro forecasting model: a callable ``(covariates, data=None) -> None``. Any plain function with this signature satisfies this Protocol structurally (for example, one that derives its `~~numpyro_forecast.models.Horizon` from the shapes via ``Horizon.from_data`` and calls the model building blocks), so nothing needs to subclass it. The parameters are positional-only so a user model's own parameter names (``cov``, ``y``, ...) stay free instead of being forced to match ``covariates``/``data``. ``ty`` checks call sites against this signature structurally (duck typing), which is the main payoff of the Protocol over a bare ``Callable`` alias. At runtime, the beartype import hook's ``isinstance`` check on a ``runtime_checkable`` Protocol only verifies that the named methods exist (Python runtime protocols never inspect signatures), so it reduces to ``callable(model)``: a model missing the ``data=None`` default still passes this check and only fails loudly at the first driver call that invokes it with ``data=None``. ForecastFn(*args, **kwargs) A closure that fits a model on a training window and forecasts its test horizon. Called by `~~numpyro_forecast.evaluate.backtest()` positionally, with ``full_covariates`` spanning the *full* window (train followed by test, i.e. ``covariates[..., t0:t2, :]``) and ``batch_size`` forwarded unchanged from ``backtest`` so a chunked closure can bound its own device memory. Returns forecast samples with the sample axis first, shape ``(num_samples, *batch, t2 - t1, obs)``. The draws may stay in host memory (e.g. via ``device="host"``): a jax Array committed to the CPU backend device or, without a CPU backend, a NumPy array. Every metric in `~~numpyro_forecast.evaluate.DEFAULT_METRICS` accepts such a ``pred`` or ``truth`` (or both), in any mix and regardless of ``batch_size``, moving a host-resident operand to device memory first where needed; draws already on-device avoid that hop for the metrics scored every window. The parameters are positional-only so a closure keeps its own parameter names. At runtime the beartype hook only checks that the value is callable (Python protocols never inspect signatures); ``ty`` checks the signature structurally at the ``backtest`` call site. Guide(*args, **kwargs) A NumPyro guide for a `ForecastModel`: a callable with the model's signature. A guide is called with the same ``(covariates, data=None)`` arguments as the model it approximates, so this Protocol has the same shape as `ForecastModel`. Both an autoguide instance (`numpyro.infer.autoguide.AutoGuide` subclasses are callables) and a hand-written guide function satisfy it structurally; nothing needs to subclass it. `~~numpyro_forecast.evaluate.backtest_vectorized()` samples an autoguide through its ``sample_posterior`` and any other guide through ``numpyro.infer.Predictive(guide, params=...)``. The same runtime caveat as `ForecastModel` applies: the beartype hook's ``isinstance`` check on a ``runtime_checkable`` Protocol reduces to ``callable(guide)`` (Python never inspects signatures at runtime), while ``ty`` checks the signature structurally at the call site. InSampleFn(*args, **kwargs) A closure that fits a model on a training window and scores its in-sample fit. Called by `~~numpyro_forecast.evaluate.backtest()` (only when ``eval_train=True``) positionally. Returns in-sample posterior-predictive samples with the sample axis first, shape ``(num_samples, *batch, t1 - t0, obs)``. The same ``batch_size``/host-offload notes as `ForecastFn` apply. Callable(*args, **kwargs) Callable(*args, **kwargs) Array() Array base class for JAX ``jax.Array`` is the public interface for instance checks and type annotation of JAX arrays and tracers. Its main applications are in instance checks and type annotations; for example:: x = jnp.arange(5) isinstance(x, jax.Array) # returns True both inside and outside traced functions. def f(x: Array) -> Array: # type annotations are valid for traced and non-traced types. return x ``jax.Array`` should not be used directly for creation of arrays; instead you should use array creation routines offered in :mod:`jax.numpy`, such as :func:`jax.numpy.array`, :func:`jax.numpy.zeros`, :func:`jax.numpy.ones`, :func:`jax.numpy.full`, :func:`jax.numpy.arange`, etc. Callable(*args, **kwargs) ## Autocorrelation Batched autocorrelation and partial autocorrelation diagnostics. acf(y: jaxtyping.Float[Array, '*batch time'] | jaxtyping.Float[ndarray, '*batch time'], max_lag: int) -> jaxtyping.Float[Array, '*batch lags'] Compute the empirical autocorrelation function up to ``max_lag``. The lag-$k$ autocorrelation of a series $y_1, \ldots, y_T$ with sample mean $\bar{y}$ is estimated with the biased normalization $$ \hat{\rho}_k = \frac{\sum_{t=k+1}^{T} (y_t - \bar{y})(y_{t-k} - \bar{y})} {\sum_{t=1}^{T} (y_t - \bar{y})^2}, $$ computed for all lags at once via the FFT of the centered series (the Wiener-Khinchin identity), so the cost is $O(T \log T)$ instead of $O(T \, k_{\max})$. The computation runs along the last axis and broadcasts over any leading batch axes. Parameters ---------- y Time series with time on the last axis, shape ``(*batch, time)``. NumPy arrays are accepted directly. max_lag Largest lag to evaluate; must satisfy ``1 <= max_lag < time``. Returns ------- Float[Array, "*batch lags"] Autocorrelations for lags ``0, 1, ..., max_lag`` (``max_lag + 1`` values; lag ``0`` is identically ``1.0``). Raises ------ ValueError If ``max_lag`` is not in ``[1, time)``. Notes ----- A constant series has zero variance, so the result is meaningless: all NaN in eager execution (a ``0 / 0``), while under `jax.jit()` the rounding of the mean can instead yield arbitrary finite values. References ---------- Adapted from `numpyro.diagnostics.autocorrelation()`, itself adapted from the Stan implementation. pacf(y: jaxtyping.Float[Array, '*batch time'] | jaxtyping.Float[ndarray, '*batch time'], max_lag: int) -> jaxtyping.Float[Array, '*batch lags'] Compute the empirical partial autocorrelation function up to ``max_lag``. The lag-$k$ partial autocorrelation is the correlation between $y_t$ and $y_{t-k}$ after removing the linear effect of the intermediate observations $y_{t-1}, \ldots, y_{t-k+1}$; it equals the last coefficient $\phi_{kk}$ of the best linear predictor of order $k$. The coefficients are obtained from the empirical autocorrelations (see `acf()`) with the Durbin-Levinson recursion $$ \phi_{kk} = \frac{\hat{\rho}_k - \sum_{j=1}^{k-1} \phi_{k-1,j}\, \hat{\rho}_{k-j}} {1 - \sum_{j=1}^{k-1} \phi_{k-1,j}\, \hat{\rho}_j}. $$ The recursion runs along the last axis and broadcasts over any leading batch axes. Parameters ---------- y Time series with time on the last axis, shape ``(*batch, time)``. NumPy arrays are accepted directly. max_lag Largest lag to evaluate; must satisfy ``1 <= max_lag < time``. Returns ------- Float[Array, "*batch lags"] Partial autocorrelations for lags ``0, 1, ..., max_lag`` (``max_lag + 1`` values; lag ``0`` is identically ``1.0``). Raises ------ ValueError If ``max_lag`` is not in ``[1, time)``. Notes ----- A constant series has zero variance, so the result is meaningless (see `acf()`); a near-deterministic series can likewise overflow the recursion (the denominator above approaches zero) and produce infinities or NaNs. References ---------- James Durbin (1960). "The Fitting of Time-Series Models". *Revue de l'Institut International de Statistique*. ## Seasonal features Fourier design matrices and seasonal tiling. fourier_features(duration: int, period: float, num_terms: int) -> jaxtyping.Float[Array, 'duration two_num_terms'] Build a Fourier seasonality design matrix. Parameters ---------- duration Number of time steps. period Seasonal period (in time steps). num_terms Number of harmonics; the output has ``2 * num_terms`` columns (sine then cosine). Returns ------- Float[Array, "duration two_num_terms"] The design matrix of shape ``(duration, 2 * num_terms)``. periodic_repeat(x: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, duration: int, *, axis: int = -1) -> jax.Array Tile a seasonal pattern to cover ``duration`` time steps. Parameters ---------- x Seasonal pattern; the repeated axis has length equal to the period. Accepts any array-like (e.g. a raw ``numpyro.sample`` draw). duration Target length along ``axis``. axis Axis to repeat along (defaults to ``-1``). Returns ------- Array ``x`` periodically repeated to length ``duration`` along ``axis``. ## Array helpers Time-axis array shaping for the train/forecast split. zero_data_like(data: jax.Array, covariates: jax.Array) -> jax.Array Return zeros shaped like ``data`` but extended to the covariate duration. Mirrors Pyro's ``zero_data``: it exposes the shape/dtype of the data over the full forecast horizon without leaking observed values into the model. The functional API exposes the equivalent value as `numpyro_forecast.models.Horizon.zero_data`. Parameters ---------- data Observed data with time at axis ``-2``, shape ``(*batch, t, obs)``. covariates Covariates with time at axis ``-2``, shape ``(*batch, duration, cov)``. Returns ------- Array Zeros of shape ``(*batch, duration, obs)``. concat_future(prefix: jax.Array, suffix: jax.Array, *, axis: int = -2) -> jax.Array Concatenate in-sample and forecast-horizon arrays along the time axis. Parameters ---------- prefix In-sample array. suffix Forecast-horizon array (same shape as ``prefix`` except along ``axis``). axis Time axis to concatenate along (defaults to ``-2``). Returns ------- Array The concatenation of ``prefix`` and ``suffix`` along ``axis``. pad_future(x: jax.Array, future: int, *, value: float = 0.0) -> jax.Array Append ``future`` rows filled with ``value`` along the time axis. The frozen-gate recipe of `numpyro_forecast.models.ssoe()`: an update gate observed in-sample (``(*batch, t_obs, obs)``) becomes a full-horizon ``xs`` leaf whose forecast rows are ``value`` (``0.0`` freezes the carry, ``1.0`` keeps an availability mask open). Zero ``future`` returns ``x`` unchanged in shape. Parameters ---------- x In-sample array with time at axis ``-2``, shape ``(*batch, t_obs, obs)``. future Number of forecast rows to append. value Fill value of the appended rows, cast to ``x.dtype``. Returns ------- Array ``x`` followed by ``future`` constant rows, shape ``(*batch, t_obs + future, obs)``. ## Datasets Example datasets used in the tutorials. load_bart_weekly() -> jaxtyping.Float[Array, 'weeks 1'] Load total weekly BART ridership (log scale) for the univariate example. Hourly counts are summed over all origin-destination pairs, aggregated into non-overlapping weeks, and log-transformed. Returns ------- Float[Array, " weeks 1"] Log weekly totals with time at axis ``-2`` and a single observation dim. load_bart_hierarchical(train_days: int = 90, test_weeks: int = 2) -> tuple[jaxtyping.Float[Array, 'origin time destin'], int, list[str]] Load the windowed hierarchical BART panel for the hierarchical example. The counts are ``log1p``-transformed and transposed to the ``(origin, time, destin)`` convention, then restricted to a ``train_days`` training window followed by a ``test_weeks`` test window. Parameters ---------- train_days Number of training days (24 hours each). test_weeks Number of test weeks (``24 * 7`` hours each). Returns ------- y : Float[Array, " origin time destin"] Log counts over the train+test window with time at axis ``-2``. split : int Index along the time axis separating train from test. stations : list[str] Station names. Raises ------ ValueError If the requested ``train_days`` + ``test_weeks`` window exceeds the available history (which would otherwise wrap a negative slice index). load_victoria_electricity() -> tuple[jaxtyping.Float[Array, 'time 1'], jaxtyping.Float[Array, 'time']] Load hourly Victoria (Australia) electricity demand and temperature. The series covers the first eight weeks of 2014, sampled hourly, from the Victoria electricity demand dataset used in the TensorFlow Probability structural-time-series case study and in Hyndman and Athanasopoulos' *Forecasting: Principles and Practice*. The original half-hourly data is downsampled to hourly by taking every other step. The values are bundled as a small CSV next to this module. Returns ------- demand : Float[Array, " time 1"] Hourly electricity demand (GW) with time at axis ``-2`` and a single observation dimension. temperature : Float[Array, " time"] Hourly temperature (degrees Celsius), aligned with ``demand``. bart_available() -> bool Return whether the BART dataset can be loaded (download succeeds). Returns ------- bool ``True`` if `load_bart_od()` loads without error. load_m5(cache_dir: str | pathlib.Path | None = None) -> numpyro_forecast.datasets.M5Data Load the M5 competition data (download and cache once, then read the files). The files come from [Nixtla's mirror](https://github.com/Nixtla/m5-forecasts) of the competition data: the daily unit sales of the 30,490 series over the training period (days 1 to 1,941) and the 28 evaluation days released after the competition, the calendar with the events and the SNAP days, the weekly shelf prices, and the official evaluation weights. The 50 MB archive is downloaded once into ``cache_dir`` (default ``~/.cache/numpyro_forecast/m5``) and checked against a pinned SHA-256 digest. Reading needs polars (``pip install numpyro_forecast[dataframes]``). Parameters ---------- cache_dir Directory that holds the archive and the extracted files; created if missing. Returns ------- M5Data Dense ``(days, series)`` sales and price arrays, the identifier table, the calendar and the official weights. M5Data(sales: jaxtyping.Float[ndarray, 'days series'], price: jaxtyping.Float[ndarray, 'days series'], keys: polars.DataFrame, calendar: polars.DataFrame, weights: polars.DataFrame) The M5 competition data as dense arrays plus the identifier and calendar tables. Attributes ---------- sales Daily unit sales ``(days, series)`` over the 1,941 training days followed by the 28 evaluation days, series in the order of the sales file. price Weekly shelf price of every series repeated over its days, ``NaN`` on the days the item was not listed. keys One row per series: ``id`` (``item_id`` and ``store_id`` joined by ``_``) and the five identifier columns ``item_id``, ``dept_id``, ``cat_id``, ``store_id``, ``state_id``. calendar One row per day with the ``date`` parsed, the Walmart week ``wm_yr_wk``, the weekday columns, the event names and types (``null`` when there is none) and the SNAP flags of the three states. weights The official evaluation weights of the 42,840 aggregates (``Level_id``, ``Agg_Level_1``, ``Agg_Level_2``, ``Dollar_Sales``, ``weight``). ## Optional dependencies Lazy imports behind pyproject extras. require(module: str, *, extra: str) -> module Import an optional dependency, or raise a targeted ``ImportError``. Optional features (dataframes, blackjax, optax) live behind ``pyproject`` extras and are never imported at package import time. This helper imports the backing module lazily at first use and, when it is missing, raises an ``ImportError`` naming the exact ``pip install`` invocation that provides it. Parameters ---------- module The importable module name (e.g. ``"pandas"``). extra The ``numpyro_forecast`` extra that installs it (e.g. ``"dataframes"``). Returns ------- ModuleType The imported module. Raises ------ ImportError If ``module`` is not importable, with an actionable install hint. ## Exceptions Package exception hierarchy raised at validation boundaries. NumpyroForecastError(message: str | None = None) -> None Base class for all deliberate ``numpyro_forecast`` errors. Subclasses may set `default_message`; instantiating without arguments then carries that message, and a positional message overrides it. Parameters ---------- message Optional explicit message; defaults to the class `default_message`. BacktestWindowError(message: str | None = None) -> None A backtest window configuration is invalid. Raised by `~~numpyro_forecast.evaluate.backtest()` when ``window_type`` and ``train_window`` disagree (``"rolling"`` without a ``train_window``, or ``"expanding"`` with one), and by `~~numpyro_forecast.evaluate.backtest_vectorized()` when ``train_window``, ``test_window``, or ``stride`` is below 1, or when the series has no room for a single window. VectorizedMetricError(message: str | None = None) -> None A metric is not vmappable in the vectorized backtest. Raised by `~~numpyro_forecast.evaluate.backtest_vectorized()` when a metric forces a host conversion (e.g. ``float(...)`` or ``numpy`` calls) under ``vmap``. Metrics must be pure JAX functions returning a scalar array; see `~~numpyro_forecast.typing.Metric`. KernelConfigError(message: str | None = None) -> None A ``contrib.blackjax`` kernel is run unbound or misconfigured. Raised by `~~numpyro_forecast.contrib.blackjax._BlackjaxKernel.init()` when the kernel was constructed with no model bound (e.g. ``BlackjaxNUTSKernel()`` instead of ``BlackjaxNUTSKernel(model)``); the fix is to pass the model as the kernel's first argument at construction time, before handing the kernel to `numpyro.infer.MCMC`. BlackJAX kernels also require ``chain_method="sequential"`` and ``num_warmup=0``; see the "Run configuration" section of `~~numpyro_forecast.contrib.blackjax.BlackjaxNUTSKernel` and its sibling kernels for why, and how misconfiguring either surfaces. CovariateDimsError(message: str | None = None) -> None Covariate dimension names are inconsistent or malformed. Raised by `~~numpyro_forecast.convert.to_datatree()` and `~~numpyro_forecast.convert.add_forecast_groups()` when ``covariate_dims`` does not name every covariates axis, or when the names passed to (or inherited by) `~~numpyro_forecast.convert.add_forecast_groups()` disagree with the dimension names already stored on the tree's ``constant_data`` covariates. MVNLayoutError(message: str | None = None) -> None A ``MultivariateNormal`` layout is unsupported for time-axis surgery. Raised by `~~numpyro_forecast.surgery.shift_loc()`, `~~numpyro_forecast.surgery.slice_time()`, and `~~numpyro_forecast.surgery.prefix_condition()` on MVN noise whose ``loc``/``covariance_matrix`` shapes do not match the supported time-leading layout. DeviceMemoryError(message: str | None = None) -> None A memory pool ran out during posterior or predictive sampling. Raised by `~~numpyro_forecast.predictive.draw_posterior()` and the predictive drivers (`~~numpyro_forecast.predictive.forecast()`, `~~numpyro_forecast.predictive.predict_in_sample()`, and everything built on them, e.g. `~~numpyro_forecast.convert.to_datatree()`) when XLA reports ``RESOURCE_EXHAUSTED``. For an accelerator OOM the message embeds the device's memory budget and the lever: the per-chunk footprint scales linearly with ``batch_size`` times the panel width, so lower (or set) ``batch_size``, free large device arrays still referenced elsewhere, and keep results off the accelerator with ``device="host"``. For a pinned host pool OOM (``Out of host memory``, reached only through an explicit ``device="pinned_host"`` or the caller's own pinned arrays) it names the pool's cap (``XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB``) and points at ``device="host"``, which lands results in pageable host memory, instead. The original XLA error is chained as ``__cause__``. HostMemoryKindError(message: str | None = None) -> None A device exposes no host memory kind for ``device="pinned_host"``. Raised by `~~numpyro_forecast._offload._host_memory_kind()`, reached from `~~numpyro_forecast.predictive.draw_posterior()`, the predictive drivers and everything built on them, when an explicit ``device="pinned_host"`` targets a device whose addressable memories include neither ``"pinned_host"`` nor ``"unpinned_host"``. The message lists the kinds the device does expose; pass a `jax.Device` or a platform name (for example ``device="cpu"``), or ``device="host"``, instead. DevicePlatformError(message: str | None = None) -> None A ``device`` platform name has no initialized JAX backend. Raised by `~~numpyro_forecast._offload._resolve_device()` (reached from every function with a ``device`` parameter) when ``device`` names a platform, e.g. ``"tpu"``, whose backend is not initialized in this process. The message lists the available platforms and how to initialize the missing one via ``jax_platforms``. ``"host"``, ``"cpu"``, ``"numpy"`` and ``"pinned_host"`` never raise this: they degrade to the NumPy path when the CPU backend is missing (see `~~numpyro_forecast.predictive.draw_posterior()`).