PyMCTarget

Note

PyMCTarget adapts a supported PyMC model to PyVBMC, including the model’s log joint, variable transforms, bounds, starting point and plausible box. Install the optional integration with python -m pip install --upgrade "pyvbmc[pymc]" on Python 3.12 or newer; see Installation. Importing pyvbmc or pyvbmc.pymc does not import PyMC or PyTensor.

PyMC adapter at a glance

Construct the model normally, then pass one adapter object directly to VBMC:

import pymc as pm
from pyvbmc import PyMCTarget, VBMC

with pm.Model() as model:
    location = pm.Normal("location", 0.0, 2.0)
    scale = pm.HalfNormal("scale", 1.0)
    pm.Normal("y", location, scale, observed=observations)

target = PyMCTarget(model, seed=7)
vbmc = VBMC(
    target,
    seed=7,
    options={"max_fun_evals": target.setup_cost + 100},
)
vp, results = vbmc.optimize()

VBMC(target) obtains x0, hard and plausible bounds, the target callable, setup observations and their function-equivalent cost from the same object. str(target) displays the layout and numerical bounds before a run starts. The read-only vbmc.target property retains the adapter for export after fitting or loading a saved run.

Explicit VBMC arguments can override the starting point, plausible bounds, observations and initialization charge. These starting points and bounds use the adapter’s flat coordinates; the adapter constructor’s start and plausible_bounds mappings use model coordinates. Hard-bound overrides must equal target.lb and target.ub in those flat coordinates, and a separate prior or log_prior is rejected because the PyMC model already includes its prior.

PyMC coordinate layout and hard bounds

Free variables follow the supplied model’s model.free_RVs order. Each variable is flattened in C order; coordinate_names uses declared PyMC coordinate labels when available. The adapter’s flatten and unflatten methods operate on PyMC value-variable names. to_model_variables and from_model_variables cross between flat VBMC points and mappings over the original model-variable names.

Variables with no transform remain in model coordinates. Variables whose coordinates all have two finite bounds also become model coordinates, with those limits as VBMC’s hard bounds. A variable whose coordinates are all one-sided keeps its PyMC transform, travels in an unbounded value coordinate such as scale_log__, and contributes the transform Jacobian to the target. This gives VBMC a fully unbounded or two-sided box while preserving the PyMC model’s density.

Explicit start= and plausible_bounds= values are mappings over every free-variable name in model coordinates. Every value must be strictly inside the recorded support. The adapter maps retained transforms to value space and applies the same final normalization used by ordinary VBMC construction.

PyMC setup and plausible-region accounting

By default, the adapter searches for a mode and builds a Laplace box of three marginal standard deviations. A non-positive-definite Hessian or an unusable marginal curvature takes its width from prior draws. A mode outside the prior’s 0.5%–99.5% location interval is reported as a diagnostic and is not moved. If the graph has no gradient, the adapter uses the model’s initial point and the 5%–95% prior quantiles instead. Supply start= or plausible_bounds= to override the corresponding automatic step.

Every prior quantity above comes from one batch of 4000 forward draws of the model’s prior. plausible_bounds= needs neither those draws nor the Hessian. start= drops the location check, so draws are then taken only for a prior-width fallback or for a model without a gradient. Both arguments are therefore the way to keep a prior that is expensive to sample forward out of setup.

The default setup_budget is 20 + 5 * D function-equivalent evaluations. It caps work before it is performed; a smaller positive budget shortens the mode search when the selected route still fits. The Hessian costs D units, and every attempted density or value-gradient call counts even when its result is rejected or non-finite. setup_cost is the actual total, while plausible_info separates target calls, Hessian calls and cost, the selected route, and any fallback diagnostics.

Finite setup observations are available in setup_evaluations and are automatically reused by VBMC(target) through the general precomputed_evaluations interface. They supplement the ordinary initial design without changing x0 or the plausible box. Importing them consumes no fresh calls; setup_cost is charged once against options["max_fun_evals"]. Consequently results["func_count"] remains the literal number of fresh calls made by VBMC, and a positive setup charge adds results["evaluation_budget"] with the total, initialization, fresh and used counts. The maximum fresh-call allowance is the configured total minus setup_cost.

Passing precomputed_evaluations=None to VBMC(target, ...) disables reuse while retaining the setup charge. Passing initialization_cost=0 retains reuse but excludes setup from that run’s budget, for example when accounting for shared preparation separately across several fits. Reusing a target does not consume its observations: each fit charges its setup cost by default. The individual fields remain available for manually constructing VBMC(target.log_joint, ...); that form requires explicit observation and cost arguments and leaves vbmc.target as None, and its runs go without the checks for a changed target described under persistence below.

Structured return to the PyMC model

Export variational draws with their model names, shapes, dimensions and coordinates, then use PyMC’s downstream functions. Use the posterior fitted with this target and inference snapshot: the dimension check cannot establish that a posterior belongs to the same model or coordinate layout. After loading a run, use loaded.target.to_arviz(loaded.vp).

posterior_data = target.to_arviz(vp, n_samples=2000)
posterior_data = pm.compute_deterministics(
    posterior_data, model=target.model, extend_dataset=True
)
predictions = pm.sample_posterior_predictive(
    posterior_data,
    model=target.model,
    var_names=["y"],
    random_seed=7,
)

The result is an ArviZ DataTree whose posterior group records parameter_space="model". Export advances vp.rng exactly as vp.sample(n_samples) does. target.model is the original prediction model, so PyMC’s ordinary pm.set_data workflow can change compatible covariates before prediction.

PyMC model-scope limits

Free variables must be continuous float64 variables with one of PyMC’s standard log, log-odds or interval transforms, or an explicitly recognized proper distribution on the whole real line. Unknown free-variable support is rejected even when no default transform is registered. A custom likelihood operation is allowed when its model graph provides a compatible density; missing derivatives select the documented setup fallbacks. This is distinct from an unknown custom free-variable distribution.

PyTensor must use floatX=float64; value variables and compiled density, gradient and Hessian outputs must also be float64. Some PyMC expressions retain a narrower constant dtype: for a uniform-only model, use explicit np.float64 bounds, for example pm.Uniform("p", np.float64(0), np.float64(1)). The adapter checks the compiled output dtype rather than casting a lower-precision result.

The adapter rejects discrete or float32 free variables, suppressed or custom transforms, interval limits that depend on another random variable, and a single variable whose coordinates mix one-sided bounds with two-sided or unbounded bounds. The log density must be finite strictly inside the hard box. The adapter returns -inf only on or outside that box, where VBMC does not evaluate. Use proper, normalized priors when interpreting the ELBO as a bound on model evidence.

A retained log transform represents support on (0, inf). A distribution whose density begins above a positive shift, such as a Wald distribution with a nonzero alpha, is therefore unsupported even when PyMC assigns it that transform. Reparameterize the variable as its distance above the shift. The setup checks may miss such a restriction; they do not guarantee that the log density is finite throughout the transformed space.

Persistence of adapted PyMC runs

Target construction freezes the model data, the numeric shared variables and NumPy arrays of the model, of its log density and of the maps between its coordinates and the model’s, and the dimensions and coordinate metadata used for inference. Later changes to the original model, or to the arrays it was built from, do not change target values or setup observations; construct a new target to refit changed data. The original model remains available for posterior prediction.

An array that the body of a pytensor.scan loop or of an OpFromGraph reads as a constant is the exception, including an array passed to a loop in non_sequences: PyTensor keeps it inside that operation, where the target does not copy it, and shares the operation between models whose arrays there are equal, so the array a target reads can be another model’s. Pass such an array to a loop in sequences or as pm.Data, and to an OpFromGraph as one of its inputs; otherwise leave it unchanged while the target is in use. For a run made with VBMC(target), a change to such an array makes optimize() refuse to run, and saving the run warns. Restoring the array allows the live run to continue unless loading has already marked the run as changed.

Saving a VBMC instance retains the adapter, its fixed inference snapshot, original prediction model, setup observations and budget charge. Loading recompiles PyTensor functions and can take several seconds. Load the file only with compatible PyMC and PyTensor versions, then use loaded.target for export and prediction. Loading a run made with VBMC(target) also checks the arrays that the target could not copy and evaluates its log density once, at the starting point, against the value recorded during setup. If the target no longer matches the run, loading warns: the run’s results remain valid, but optimize refuses to continue it; build a new target and run to fit the changed data.

class pyvbmc.pymc.PyMCTarget(model, *, plausible_bounds=None, start=None, seed=None, setup_budget=None)

Represent a supported PyMC model as a flat PyVBMC target.

The adapter snapshots the model inputs at construction, maps supported constrained variables to coordinates accepted by PyVBMC, and prepares a finite starting point and plausible box. Setup density evaluations are retained for reuse by pyvbmc.VBMC.

Parameters:

modelpymc.Model

Model whose free variables define the target, in model.free_RVs order.

plausible_boundsmapping, optional

Mapping from every free-variable name to (lower, upper) values in model coordinates. Values must broadcast to the variable shape and be strictly inside its support. By default a Laplace box with prior-width fallbacks is constructed; supplying bounds instead needs neither the Hessian nor the 4000 prior draws that box can require, which matters when the prior is expensive to sample forward.

startmapping, optional

Starting point in model coordinates. It must contain every free variable and lie strictly inside its support. Supplying it skips the mode search and the automatic box’s location check against prior draws, which are then drawn only for a prior-width fallback or for a model without a gradient.

seedNone, int, SeedSequence or numpy.random.Generator, optional

Random generator or seed for prior fallback draws and stochastic initval="prior" strategies. A generator is used as is. With None, a generator is derived from NumPy’s legacy global state.

setup_budgetint, optional

Positive function-equivalent setup allowance. The default is 20 + 5 * D. Actual work, including rejected search trials and a D-unit Hessian evaluation, determines setup_cost.

Attributes:

Dint

Dimension of the flat PyVBMC coordinates.

names, value_names, coordinate_nameslist of str

Model variable names, PyVBMC value-variable names, and flattened coordinate labels.

shapeslist of tuple

Model-variable shapes in layout order.

sizeslist of int

Flattened sizes in layout order.

keptdict

Model-variable names mapped to the class names of transforms retained in PyVBMC coordinates.

supportdict

Model-variable names mapped to (lower, upper) support arrays.

x0, lb, ub, plb, pubnumpy.ndarray

Float64 row vectors of shape (1, D).

setup_evaluationstuple of numpy.ndarray

Finite setup observations (X, y) in call order.

plausible_infodict

Setup route, diagnostics, and exact cost accounting.

modelpymc.Model

The original model, retained for posterior prediction. Numerical target calls use a private frozen snapshot.

Raises:

ImportError

If the PyMC extra is absent or the installed PyMC API is incompatible.

UnsupportedModel

If a free variable, transform, support, or dtype is unsupported.

ValueError

If setup arguments, bounds, or the final starting density are invalid.

Notes

A retained log transform represents support on (0, inf). A distribution whose density begins above a positive shift, such as a Wald distribution with a nonzero alpha, is unsupported; reparameterize the variable as its distance above the shift. The setup checks may miss such a restriction; they do not guarantee that the log density is finite throughout the transformed space.

flatten(values)

Flatten a mapping over value-variable names in C order.

Parameters:

valuesmapping

One array of the declared shape for every name in value_names.

Returns:

xnumpy.ndarray

Float64 vector of shape (D,).

from_model_variables(point)

Map one model-variable point to a flat PyVBMC vector.

Parameters:

pointmapping

Every model-variable name mapped to one finite value in model coordinates. Values may be broadcast to the variable’s declared shape and must lie strictly inside its support.

Returns:

xnumpy.ndarray

Float64 vector of shape (D,) in the adapter’s flat PyVBMC coordinates. Retained PyMC transforms are applied.

Raises:

ValueError

If point is not a mapping over exactly the model-variable names, a value cannot be broadcast to its declared shape, or a value is non-finite or outside its support.

log_joint(x)

Evaluate the model log joint, including retained Jacobians.

A point on or outside a hard bound returns -inf. A non-finite density strictly inside the box raises a ValueError that identifies the model variables and their values.

log_joint_no_jacobian(x)

Evaluate the model-coordinate log density without Jacobians.

The input remains in the adapter’s flat PyVBMC coordinates. This value omits the Jacobian terms of any retained PyMC transforms.

Parameters:

xarray-like

One point with shape (D,) or (1, D) in flat PyVBMC coordinates.

Returns:

float

The model-coordinate log density. A point on or outside a hard bound returns -inf.

Raises:

ValueError

If x has the wrong shape, or if the compiled density is non-finite strictly inside the hard bounds.

property setup_cost

Actual function-equivalent setup cost.

to_arviz(vp, n_samples=1000)

Export variational draws in the PyMC model’s coordinates.

Parameters:

vpVariationalPosterior

Posterior fitted with this target’s frozen model and data, in its flat coordinate layout. Only dimension compatibility is checked. After loading a run, use loaded.target.to_arviz(loaded.vp).

n_samplesint, optional

Positive number of independent draws, default 1000, as any scalar that holds a whole number (see VariationalPosterior.sample); a boolean is refused.

Returns:

dataxarray.DataTree

One posterior group with the model’s variable names, shapes, dimensions, and coordinates.

Raises:

ImportError

If the ArviZ data API is unavailable.

ValueError

If the sample count, posterior dimension, or stored layout is invalid. Validation errors leave vp.rng unchanged.

to_model_variables(X)

Map flat PyVBMC draws to arrays over the model’s variables.

Parameters:

Xnumpy.ndarray

Float64-compatible array of shape (n, D).

Returns:

valuesdict

Variable names mapped to arrays of shape (n, *shape). Every value is finite and strictly inside its support. A retained transform’s backward map can land on a finite support bound by rounding, and such a coordinate is returned as the adjacent representable value inside the support.

Raises:

ValueError

If X has the wrong shape, or a mapped value is non-finite or lies beyond a support bound.

unflatten(x)

Convert one flat PyVBMC point to a value-variable mapping.

Parameters:

xarray-like

One point with shape (D,) or (1, D) in the adapter’s flat PyVBMC coordinates.

Returns:

valuesdict

PyMC value-variable names mapped to float64 arrays with the variables’ declared shapes. Retained PyMC transforms are not inverted.

Raises:

ValueError

If x does not have shape (D,) or (1, D).

exception pyvbmc.pymc.UnsupportedModel

Raised when a PyMC model is outside the adapter’s supported scope.