What bayesmith does and why

bayesmith compiles a Bayesian model, written as a graph of operators, into an inference procedure that can be audited. It finds and checks the structure of the graph, solves exactly the parts that admit a closed form, hands the remainder to a numerical backend, and returns results that carry their conditions, diagnostics and provenance. It sits between the model declaration and the sampler. It is not a probabilistic programming language and not a collection of samplers.

One model, before an algorithm

Let \mathbf z collect all unknowns, \mathbf d the observed data, and \mathbf c the fixed inputs. A model combines probability factors with deterministic calculations. After evaluating the deterministic nodes, its posterior has the general form

Each factor depends on the variables it names: a prior or conditional prior, an observation likelihood, or an explicitly declared graph-level density term. A deterministic node computes s_k=f_k(\mathrm{pa}(s_k)); its value is substituted into downstream factors. It contributes neither a free sampled coordinate nor a separate probability factor in this expression. The factors need not be Gaussian, the functions need not be linear, and a vector-valued likelihood can describe correlated data. Which parts admit a supported inference route is a separate question.

For exactly this drawing, with \mathbf z=(\theta,x,\psi), the joint density is

The drawing has four probability factors, two deterministic calculations and one fixed input. It assumes independent root priors for θ and ψ; other models can connect those variables or add a joint prior factor. Each symbol may represent a scalar, vector or field. No Gaussian law, linear response or repetition axis is assumed.

A sample node can have parents

A hyperparameter is a modeling role, not a separate node type. When it is unknown, declare it with sample and give it a hyperprior; when fixed, use a constant; when derived from other nodes, use det. A downstream sample can use any of these as parents to construct its conditional distribution. Parameters can control a mean, covariance, scale or shape. Inference conditions on the data and infers the unknown hyperparameters together with the other latents.

sample and observe both declare probabilistic nodes; the latter attaches observed values. During trace, these calls record the model rather than draw a simulated dataset. plate annotates repeated structure, and joint_prior supplies an additional density factor under its own contract; neither is an extra node class. See Graph and hierarchy for these declarations and their boundaries.

Why structure should choose the inference

Holding the other unknowns fixed, a block can have a tractable conditional posterior even when the complete model does not. For example, a Gaussian conditional prior and Gaussian observation factors with means affine in the block and covariances independent of it give a Gaussian conditional posterior, provided no other factor involving that block breaks this form. bayesmith checks the supported structural and numerical conditions before assigning an exact route. A log-space route additionally requires the transformed prior and likelihood to have the corresponding Gaussian form; positivity or an exponential response alone is insufficient. Approximate noise transformations retain their approximation status.

In common practice this choice is made by hand. Applying a general sampler to the whole model can ignore exploitable structure; hand-written conditional updates can silently become invalid when the model changes. bayesmith derives a plan from the graph and the statistical task, checks the conditions its routes require, and records their evidence. A chain using exact conditional updates can still be correlated and must reach its target distribution.

A concrete example: infer the hyperparameters too

The runnable example below instantiates the same diagram. θ is theta, η is scale, ψ is sigma, c is the fixed matrix A, and s is signal. The prior width of x and the observation noise are both inferred:

Normal distributions here use variance or covariance as their second mathematical parameter; the Python constructors use standard deviation. The defining dependency is visible in the declaration:

theta = bs.sample("theta", lambda: dist.Normal(0.0, 0.5))
scale = bs.det("scale", jnp.exp, theta)
x = bs.sample(
    "x", lambda sd: dist.Normal(jnp.zeros(2), sd).to_event(1), scale
)

The final argument scale is a parent node reference. Its current value is passed to sd when the distribution is evaluated. This declares x | scale, while the deterministic node computes scale = exp(theta). The observation likewise has two parents: signal and sigma.

"""Run with JAX_ENABLE_X64=1 python overview_hierarchy.py."""

import jax.numpy as jnp
import numpy as np
import numpyro.distributions as dist

import bayesmith as bs

A = jnp.array([[1.0, 0.3], [0.2, 1.0], [0.5, -0.4], [1.0, 1.0]])
DATA = jnp.array([0.4, 1.1, -0.2, 1.5])


def model(data):
    # A hyperparameter is an ordinary latent sample node.
    theta = bs.sample("theta", lambda: dist.Normal(0.0, 0.5))
    scale = bs.det("scale", jnp.exp, theta)

    # The distribution is constructed from the parent's current value.
    x = bs.sample("x", lambda sd: dist.Normal(jnp.zeros(2), sd).to_event(1), scale)
    design = bs.const("A", A)
    signal = bs.det("signal", lambda matrix, value: matrix @ value, design, x)

    # An observed distribution can have uncertain parameters too.
    sigma = bs.sample("sigma", lambda: dist.LogNormal(-1.0, 0.3))
    bs.observe("d", lambda mean, sd: dist.Normal(mean, sd), signal, sigma, obs=data)


def normal_log_density(value, mean, sd):
    """Independent Gaussian formula, including its scale normalization."""
    return np.sum(
        -0.5 * ((value - mean) / sd) ** 2 - np.log(sd) - 0.5 * np.log(2 * np.pi)
    )


def verify_density():
    """Check the displayed factorization at 27 points, not one fitted mode."""
    graph = bs.trace(model, DATA)
    errors = []
    for theta in (-0.7, 0.0, 0.5):
        for sigma in (0.2, 0.5, 1.1):
            for x in (np.zeros(2), np.array([-0.4, 0.7]), np.array([1.0, -0.3])):
                expected = (
                    normal_log_density(theta, 0.0, 0.5)
                    + normal_log_density(x, 0.0, np.exp(theta))
                    + normal_log_density(np.log(sigma), -1.0, 0.3)
                    - np.log(sigma)
                    + normal_log_density(np.asarray(DATA), np.asarray(A) @ x, sigma)
                )
                actual = float(
                    bs.log_joint(
                        graph,
                        {
                            "theta": jnp.asarray(theta),
                            "x": jnp.asarray(x),
                            "sigma": jnp.asarray(sigma),
                        },
                    )
                )
                np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12)
                errors.append(abs(actual - expected))
    return max(errors)


if __name__ == "__main__":
    print(bs.compile(bs.trace(model, DATA)))
    print("27 density checks passed; maximum absolute error:", verify_density())

Compile the graph before sampling. The actual recorded plan is

block 0  {x}              GCR exact         structurally certified; numerical check 3 scales x 3 at-points (max 0.00e+00)
         kappa in [2.2003166, 485.25358] -> tol=2.0607782e-06, guard hoisted out of the
         sweep
         exact block ['x']: sigma does not move with the block (0.000e+00 relative,
         rtol=1.0e-08), so the covariance is fixed and one Wiener solve is exact; the
         linearity check held ['sigma'] at their prior centre rather than drawing them
         -- they are not Gaussian, so their prior has no draw this module will take; the
         kappa sweep held ['sigma'] at their prior centre rather than moving them --
         they are not Gaussian, so this module has no prior width to sweep them by, and
         the interval is not an interval in them
block 1  {sigma, theta}   NUTS
         outside the exact block ['x']: sampled by NUTS, and the exact block is solved
         CONDITIONAL on them, once per sweep
execution: HMCGibbs(inner=NUTS, gibbs_sites=['x']); noise_std rebuilt every sweep

With θ and σ fixed, the prior on x is Gaussian, the forward response is affine in x and the Gaussian noise covariance is fixed with respect to x. The plan therefore gives x an exact conditional update and sends θ and σ to NUTS. The derived scale and signal and fixed A add no sampled coordinates. They still affect the density and the inference: changing θ changes the prior for x, and changing σ changes its likelihood. Each sweep uses the current hyperparameter values.

Simpler structure and a rejected claim

The mixed example illustrates a different dependency: x and ν have independent priors, and ν scales the forward response. The curved counterexample has only x and declares exp(A @ x) affine in x. The source below declares these two alternative models.

"""The plans on the Overview page. Run with: JAX_ENABLE_X64=1 python overview_plan.py"""

import jax.numpy as jnp
import numpyro.distributions as dist

import bayesmith as bs

A = jnp.array([[1.0, 0.3], [0.2, 1.0], [0.5, -0.4], [1.0, 1.0]])
DATA = jnp.array([0.4, 1.1, -0.2, 1.5])


def mixed(data):
    """Affine in x once nu is fixed; not affine in nu. Nothing is declared."""
    x = bs.sample("x", lambda: dist.Normal(jnp.zeros(2), 2.0).to_event(1))
    nu = bs.sample("nu", lambda: dist.LogNormal(0.0, 0.5))
    mu = bs.det("mu", lambda x_, nu_: (A @ x_) * jnp.exp(0.1 * nu_), x, nu)
    bs.observe("d", lambda m: dist.Normal(m, 0.5), mu, obs=data)


def curved(data):
    """Declared affine in x, but exp(A x) is not."""
    x = bs.sample("x", lambda: dist.Normal(jnp.zeros(2), 2.0).to_event(1))
    mu = bs.det("mu", lambda x_: jnp.exp(A @ x_), x, linear_in=("x",))
    bs.observe("d", lambda m: dist.Normal(m, 0.5), mu, obs=data)


if __name__ == "__main__":
    for model in (mixed, curved):
        print(f"# {model.__name__}")
        print(bs.compile(bs.trace(model, DATA)))

block 0  {x}              GCR exact         structurally certified; numerical check 3 scales x 3 at-points (max 0.00e+00)
         kappa=72.652432 -> tol=1.3764164e-05, guard hoisted out of the sweep
         exact block ['x']: sigma does not move with the block (0.000e+00 relative,
         rtol=1.0e-08), so the covariance is fixed and one Wiener solve is exact; the
         linearity check held ['nu'] at their prior centre rather than drawing them --
         they are not Gaussian, so their prior has no draw this module will take; the
         kappa sweep held ['nu'] at their prior centre rather than moving them -- they
         are not Gaussian, so this module has no prior width to sweep them by, and the
         interval is not an interval in them
block 1  {nu}             NUTS
         outside the exact block ['x']: sampled by NUTS, and the exact block is solved
         CONDITIONAL on them, once per sweep
execution: HMCGibbs(inner=NUTS, gibbs_sites=['x']); noise_std rebuilt every sweep

block 0  {x}              NUTS
         'x': latent 'x' is declared linear, but the prediction is not affine in it:
         departure from its own linearization is too large at [0.001, 1.0, 1000.0] times
         each latent's declared prior width, evaluated at the caller's own `at`. Two
         criteria, per element, over the elements clearing the per-element roundoff
         floor -- either alone is a refusal: the relative departure against
         rtol=2.22e-12, and the departure in units of the noise sigma against
         weighted_rtol=1.00e-03. Worst per probe (relative | sigma-weighted): 0.001x ->
         1.01e-03 | 4.06e-06, 1x -> 1.30e+00 | 1.84e+01, 1000x -> nan | nan. An
         `unresolved:` entry there is a departure the roundoff floor declined to judge,
         not one measured as zero -- it is not what refused this claim, and `with
         jax.enable_x64(True):` is what settles it. Either drop the linear_in
         declaration, or re-parameterize so the model really is affine there. For a
         group that is only pairwise affine, split it into separate blocks and
         alternate.
execution: NUTS

What it does across a workflow

The design follows one model through six steps. The status column summarizes capabilities and release status.

Through the typed task interface, a run returns a Result together with the plan that produced it and a run record: seed, precision, backend, budget, termination reason and whether the answer is exact or Monte Carlo. In the task protocol, a question the available routes cannot answer returns a Refusal naming the failed premise and the measurements behind it. Faults in the model itself and failed convergence are raised as exceptions.

How it differs from the libraries it uses

Inference libraries already provide algorithms for particular structures, such as latent-Gaussian or Laplace-marginalized models. bayesmith's own job is to find and verify structure in the whole graph, combine exact and numerical routes across its parts, and state why the combination is valid for the task being asked.

Tutorials and examples

One numbered learning path, with runnable source and recorded results. Start with Example 01 for the basic API, explore model structure in 02–04, then follow a result through tasks and checks in 05–06.

Where to read next

This documentation describes version 0.10.0. See capabilities and release status.