API Reference

Model building blocks

Plain model functions that register the train/forecast sites for you.

models.Horizon: The train/forecast split for a single model call.
models.Transition: (carry, x_t) -> dist_t: the distribution of the next latent given the carry.
models.Advance: (carry, z_t, x_t) -> carry: the next carry from the current carry, the
models.innovations(): Sample conditionally iid per-step innovations over the full horizon.
models.markov_series(): Sample a Markov (state-space) latent over the full horizon.
models.ssoe(): Run a single-source-of-error recursion over the full horizon.
models.SSOEMean: (carry, x_t) -> mu_t: the one-step-ahead mean of the current row, shape
models.SSOEUpdate: (carry, y_t, eps_t, x_t) -> carry: the next carry from the row’s value and
models.SSOEResult: The means and sampled future values produced by ssoe().
models.predict(): Register the observation and forecast sites for the model.
models.PlateName: Names of the plates the building blocks open, as str members.

Vector autoregression

VAR components that compose with ssoe and markov_series: conditional mean, step factory, companion matrix, impulse responses.

var.var_mean(): Compute the VAR conditional mean of the next row from a lag window.
var.var_step(): Build the ~~numpyro_forecast.models.ssoe() mean and update of a VAR.
var.companion_matrix(): Stack the VAR coefficients into the companion form of a VAR(1).
var.impulse_response(): Compute the impulse responses (moving-average coefficients) of a VAR.

Priors

Shrinkage prior moments for coefficient arrays.

priors.minnesota_prior(): Return the Minnesota (Litterman) prior moments for VAR coefficients.

Distribution surgery

Time-axis operations on observation distributions, extensible via singledispatch.

surgery.shift_loc(): Re-center a zero-centered noise distribution at loc.
surgery.slice_time(): Slice an elementwise distribution along the time axis -2.
surgery.prefix_condition(): Condition a (t+f)-length distribution on a t-length data prefix.
surgery.register_elementwise(): Declare a distribution family elementwise (usable as a decorator).

Reparameterization

Time-axis reparameterization of in-sample latents (Haar / DCT), after Pyro’s time_reparam.

reparam.time_reparam(): Reparameterize every in-sample time latent of model along the time axis.
reparam.TimeTransform: The time-axis transform applied by time_reparam(): "haar" or "dct".

Producing draws

Drawing posterior samples and generating forecasts and in-sample predictions.

predictive.draw_posterior(): Draw num_samples posterior samples of the latent sites from a fitted guide.
predictive.forecast(): Sample forecasts for the steps in [t, duration) from a posterior.
predictive.predict_in_sample(): Sample the in-sample posterior predictive of the obs site.

Backtesting & evaluation

Rolling-window backtesting and forecast metrics.

evaluate.backtest(): Backtest a forecasting model on a moving window of (train, test) data.
evaluate.backtest_vectorized(): Rolling-window backtest with all windows fitted in one vmapped SVI run.
evaluate.BacktestResult: Per-window result of a backtest() run.
evaluate.VectorizedBacktestResult: Result of a backtest_vectorized() run (all windows at once).
evaluate.WindowType: Backtest windowing strategy: an expanding (t0=0) or fixed-size rolling window.
evaluate.DEFAULT_METRICS: Default metrics used by backtest() and backtest_vectorized().
evaluate.evaluate_forecast(): Evaluate forecast samples against ground truth for several metrics at once.
evaluate.results_to_dataframe(): Flatten backtest results into a tidy one-row-per-window DataFrame.
evaluate.eval_crps(): Empirical CRPS averaged over all data elements.
evaluate.eval_mae(): Mean absolute error using the forecast sample median as point estimate.
evaluate.eval_rmse(): Root mean squared error using the forecast sample mean as point estimate.
evaluate.eval_coverage(): Empirical coverage of the central alpha prediction interval.
metrics.crps_empirical(): Compute the empirical Continuous Ranked Probability Score (CRPS).
metrics.eval_pinball(): Mean pinball (quantile) loss of the forecast quantile.
metrics.eval_interval_score(): Mean Winkler interval score for the central alpha prediction interval.
metrics.make_mase(): Build a Mean Absolute Scaled Error metric scaled by train_data.

ArviZ export

Convert posteriors into ArviZ-schema xarray DataTrees for diagnostics and plotting.

convert.to_datatree(): Convert an already-drawn posterior into an ArviZ-schema xarray.DataTree.
convert.add_forecast_groups(): Attach out-of-sample forecast groups to a copy of tree.
convert.predictions_to_datatree(): Pack prediction draws into a DataTree laid out for per-series plot_lm faceting.

Extensions (contrib)

Optional backends behind pyproject extras (never imported by default).

contrib.blackjax.BlackjaxNUTSKernel: BlackJAX NUTS with Stan-style window adaptation.
contrib.blackjax.BlackjaxMCLMCKernel: BlackJAX Microcanonical Langevin Monte Carlo (MCLMC).
contrib.blackjax.BlackjaxCustomKernel: Adapt an arbitrary BlackJAX sampler via a user-supplied build_fn.
contrib.blackjax.PathfinderFit: The result of fitting a forecasting model with BlackJAX Pathfinder.
contrib.blackjax.fit_pathfinder(): Fit a forecasting model with BlackJAX Pathfinder variational inference.
contrib.blackjax.pathfinder_samples(): Draw num_samples posterior samples from a fitted Pathfinder approximation.
contrib.blackjax.MultiPathfinderFit: The result of fitting a forecasting model with multi-path BlackJAX Pathfinder.
contrib.blackjax.fit_multipathfinder(): Fit a forecasting model with multi-path BlackJAX Pathfinder and PSIS resampling.
contrib.blackjax.multipathfinder_samples(): Draw num_samples posterior samples from a fitted multipath Pathfinder fit.

Typing

Public type contracts.

typing.ForecastModel: A NumPyro forecasting model: a callable (covariates, data=None) -> None.
typing.ForecastFn: A closure that fits a model on a training window and forecasts its test horizon.
typing.Guide: A NumPyro guide for a ForecastModel: a callable with the model’s signature.
typing.InSampleFn: A closure that fits a model on a training window and scores its in-sample fit.
typing.Metric: A metric maps (pred, truth) forecast samples and ground truth to a scalar array.
typing.ModelFactory: A zero-argument callable returning a fresh ForecastModel instance.
typing.Array: A JAX array (alias of jax.Array).
typing.BlackjaxBuildFn: A blackjax sampler build function (rng_key, logdensity_fn, position, num_warmup).

Autocorrelation

Batched autocorrelation and partial autocorrelation diagnostics.

acf.acf(): Compute the empirical autocorrelation function up to max_lag.
acf.pacf(): Compute the empirical partial autocorrelation function up to max_lag.

Seasonal features

Fourier design matrices and seasonal tiling.

features.fourier_features(): Build a Fourier seasonality design matrix.
features.periodic_repeat(): Tile a seasonal pattern to cover duration time steps.

Array helpers

Time-axis array shaping for the train/forecast split.

arrays.zero_data_like(): Return zeros shaped like data but extended to the covariate duration.
arrays.concat_future(): Concatenate in-sample and forecast-horizon arrays along the time axis.
arrays.pad_future(): Append future rows filled with value along the time axis.

Datasets

Example datasets used in the tutorials.

datasets.load_bart_weekly(): Load total weekly BART ridership (log scale) for the univariate example.
datasets.load_bart_hierarchical(): Load the windowed hierarchical BART panel for the hierarchical example.
datasets.load_victoria_electricity(): Load hourly Victoria (Australia) electricity demand and temperature.
datasets.bart_available(): Return whether the BART dataset can be loaded (download succeeds).

Optional dependencies

Lazy imports behind pyproject extras.

optional.require(): Import an optional dependency, or raise a targeted ImportError.

Exceptions

Package exception hierarchy raised at validation boundaries.

exceptions.NumpyroForecastError: Base class for all deliberate numpyro_forecast errors.
exceptions.BacktestWindowError: A backtest window configuration is invalid.
exceptions.VectorizedMetricError: A metric is not vmappable in the vectorized backtest.
exceptions.KernelConfigError: A contrib.blackjax kernel is run unbound or misconfigured.
exceptions.CovariateDimsError: Covariate dimension names are inconsistent or malformed.
exceptions.MVNLayoutError: A MultivariateNormal layout is unsupported for time-axis surgery.
exceptions.DeviceMemoryError: A memory pool ran out during posterior or predictive sampling.
exceptions.HostMemoryKindError: A device exposes no host memory kind for device="pinned_host".
exceptions.DevicePlatformError: A device platform name has no initialized JAX backend.