Probabilistic Programming

Reactant.jl compiles ordinary Julia code through MLIR, running it on CPU, GPU, or TPU without rewriting. You write a normal Julia function, wrap a call to it in @compile, and Reactant traces the function, lowers it through MLIR, and hands you back a callable compiled program. Array inputs are staged through ConcreteRArray (created with Reactant.to_rarray) so they live on the target device.

Reactant.ProbProg is the Julia front-end for the impulse dialect, implemented across Enzyme (dialect definition and inference materialization passes) and Enzyme-JAX (backend-specific lowering). The impulse dialect provides high-level MLIR ops for describing probabilistic modeling and inference, materializes inference computation through compiler passes, and applies general-purpose and probabilistic-programming-specific optimizations during lowering.

optimize = :probprog opt-in required

For now, @compile needs an explicit optimize = :probprog argument on probabilistic programs to enable the impulse-specific MLIR passes (you'll see this in every @compile call below). Merging those passes into the default @compile pipeline is work in progress; once it lands, the explicit opt-in will no longer be required.

Next, we walk through two operating modes of Reactant.ProbProg: a trace-based mode built around a generative function, and a custom log-density mode that takes a custom log-density function.

Trace-based mode

We describe a Bayesian linear regression question:

$$\begin{array}{rl}\text{slope}& \sim \mathcal{N}(0,2)\\ \text{intercept}& \sim \mathcal{N}(0,10)\\ y_i\mid \text{slope},\text{intercept}& \sim \mathcal{N}(\text{slope}\cdot x_i+\text{intercept},1)\end{array}$$

Both regression coefficients are given Gaussian priors, tighter on slope (standard deviation 2) and looser on intercept (standard deviation 10). Each observation y_i is then drawn from a Gaussian centered on the fitted value slope · x_i + intercept with fixed noise (standard deviation 1).

The data

We start by drawing synthetic data using a known slope and intercept of -2 and 10 respectively.

true_slope, true_intercept = -2.0, 10.0

xs = collect(Float64, 1:10)
ys = (true_slope .* xs) .+ true_intercept .+ randn(length(xs))
(xs, ys)

([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], [7.291565677973274, 6.141286598450446, 2.6394702721493166, 2.4885796184136235, 0.6299275833857677, -1.8192611855435525, -5.160684756507587, -6.064152948401424, -8.49825800409541, -11.307936349145626])

Describing the Model

We describe this model in Reactant.ProbProg as follows:

using Reactant: ProbProg

function model(rng, xs)
    _, slope = ProbProg.sample(
        rng, ProbProg.Normal(0.0, 2.0, (1,)); symbol=:slope,
    )
    _, intercept = ProbProg.sample(
        rng, ProbProg.Normal(0.0, 10.0, (1,)); symbol=:intercept,
    )
    _, ys = ProbProg.sample(
        rng,
        ProbProg.Normal(slope .* xs .+ intercept, 1.0, (length(xs),));
        symbol=:ys,
    )
    return ys
end

model (generic function with 1 method)

Each random choice is introduced by a ProbProg.sample(rng, dist; symbol=...) call that takes a random number generator (RNG) and a distribution function. The symbol keyword names the sample site used for conditioning and specifying parameters to infer.

As a calling convention, ProbProg.sample returns (rng, value); the first element (omitted with _ above) is the updated RNG. In the current implementation rng is a traced ReactantRNG whose state corresponds to a tensor<2xui64> RNG state in the generated MLIR. We don't thread it through manually because Reactant tracing handles the input/output threading at the IR level, and ReactantRNG's internal state is updated via Julia mutability (see here for details).

Describing Inference

We condition on the observed ys with a Constraint object:

obs = ProbProg.Constraint(:ys => ys)

Reactant.ProbProg.Constraint with 1 entry:
  Address([:ys]) => [7.29157, 6.14129, 2.63947, 2.48858, 0.629928, -1.81926, -5…

The current implementation requires a bit of boilerplate to flatten the Constraint into a tensor representation and to extract its address set before passing them to the @compile'd function below (see Traces and constrained inference for details):

obs_tensor = ProbProg.flatten_constraint(obs)

1×10 ConcretePJRTArray{Float64,2}:
 7.29157  6.14129  2.63947  2.48858  …  -6.06415  -8.49826  -11.3079

constrained_addresses = ProbProg.extract_addresses(obs)

Set{Reactant.ProbProg.Address} with 1 element:
  Reactant.ProbProg.Address([:ys])

We then specify what parameters to infer:

selection = ProbProg.select(
    ProbProg.Address(:slope),
    ProbProg.Address(:intercept),
)

OrderedCollections.OrderedSet{Reactant.ProbProg.Address} with 2 elements:
  Reactant.ProbProg.Address([:slope])
  Reactant.ProbProg.Address([:intercept])

We express inference in a single function that conditions the model on the constraint with generate and then runs NUTS over the selected sites with mcmc.

function infer(rng, xs, obs_tensor, step_size, inverse_mass_matrix)
    trace, = ProbProg.generate(
        rng, obs_tensor, model, xs; constrained_addresses,
    )
    trace, = ProbProg.mcmc(
        rng, trace, model, xs;
        selection, algorithm=:NUTS,
        step_size, inverse_mass_matrix,
        num_warmup=200, num_samples=500,
    )
    return trace
end

infer (generic function with 1 method)

The returned trace contains the sampling result as a 2D tensor: each row is the concatenation of all selected sites' flattened values for one post-warmup sample. (We will show a possible trace for this example problem below.)

Compiling with @compile

We compile infer with Reactant's @compile for compiler-optimized probabilistic inference:

rng                 = ReactantRNG()
step_size           = Reactant.ConcreteRNumber(0.1)
inverse_mass_matrix = Reactant.ConcreteRArray([1.0 0.0; 0.0 1.0])

compiled_fn = @compile optimize=:probprog infer(
    rng, xs, obs_tensor, step_size, inverse_mass_matrix,
)

Reactant compiled function infer (with tag ##infer_reactant#1121)

Defaults

It is often sufficient to start with step_size = 1.0 and an identity inverse_mass_matrix. With the default adapt_step_size = true and adapt_mass_matrix = true, mcmc adaptively selects appropriate values during the warmup iterations.

The compiled_fn is a callable object that takes the same arguments as infer and returns the inference result. We can execute the compiled inference program any number of times by calling it:

trace_tensor = compiled_fn(rng, xs, obs_tensor, step_size, inverse_mass_matrix)

500×2 ConcretePJRTArray{Float64,2}:
 -2.10869  10.7667
 -2.06337   9.81463
 -1.96603   9.39184
 -2.07439  10.2791
 -2.0589    9.67148
 -1.99333   9.63661
 -1.97273   9.5606
 -2.09112  10.7461
 -2.16264  11.0108
 -2.08233  10.1085
  ⋮        
 -2.06829   9.70998
 -1.98706   9.68992
 -2.09529   9.8552
 -1.99049   9.87547
 -1.971     9.87203
 -2.12101  10.39
 -1.89742   9.1124
 -2.04288   9.42008
 -2.04702   9.55026

In this array, each row is one post-warmup sample, and each column is a parameter. The columns follow the order that we passed to ProbProg.select above, i.e., :slope first, then :intercept.

           :slope   :intercept
sample 1:    ...       ...    
sample 2:    ...       ...    
   ⋮          ⋮         ⋮     
sample N:    ...       ...

The posterior means obtained are:

mean(trace_tensor, dims=1)

1×2 ConcretePJRTArray{Float64,2}:
 -2.0499  9.88584

We can see that NUTS recovers the values that were used to generate the data.

Custom logpdf mode

In larger applications, it is often infeasible to express the model in a PPL modeling language as we showed in the trace-based mode above. We can use Reactant.ProbProg to compile and run its inference algorithms directly on a hand-written log-density function via the custom logpdf mode.

For example, we can write the log-density function of the previous Bayesian linear regression model directly:

function logdensity(θ, xs, ys)
    X = hcat(xs, ones(length(xs)))
    residuals = ys .- X * θ
    pr = -0.5 * sum(θ .^ 2 ./ [4.0, 100.0])
    ll = -0.5 * sum(residuals .^ 2)
    return ll + pr
end

logdensity (generic function with 1 method)

We pass logdensity to the mcmc_logpdf interface along with an initial position vector (the parameter values the chain starts from):

function infer_logpdf(rng, θ0, xs, ys, step_size, inverse_mass_matrix)
    trace, = ProbProg.mcmc_logpdf(rng, logdensity, θ0, xs, ys;
        algorithm=:NUTS,
        step_size, inverse_mass_matrix,
        num_warmup=200, num_samples=500,
    )
    return trace
end

θ0 = Reactant.to_rarray(reshape([0.0, 0.0], 1, 2))
compiled_logpdf = @compile optimize=:probprog infer_logpdf(
    rng, θ0, xs, ys, step_size, inverse_mass_matrix,
)
trace = compiled_logpdf(rng, θ0, xs, ys, step_size, inverse_mass_matrix)

500×2 ConcretePJRTArray{Float64,2}:
 -2.04357  10.0199
 -2.21348  10.76
 -2.13843  10.328
 -2.16822  10.5459
 -2.18175  10.6627
 -1.9106    8.65892
 -2.18314  10.7746
 -1.797     8.2119
 -2.13959  10.0646
 -2.05409   9.64582
  ⋮        
 -2.11362   9.99909
 -2.05747   9.93262
 -1.95564   9.44108
 -1.79311   8.68817
 -1.99031   9.47834
 -2.11341  10.2017
 -1.92273   9.5201
 -2.03557   9.60742
 -2.15958  10.9585

We get similar inference results

(
    posterior_mean_slope     = mean(trace[:, 1]),
    posterior_mean_intercept = mean(trace[:, 2]),
)

(posterior_mean_slope = -2.039233042923542, posterior_mean_intercept = 9.822228167820585)

NUTS recovers both posterior means here too.

More Explanations

• Sampling and distributions — semantics of sample, the built-in Distribution hierarchy, custom samplers with user-supplied logpdf, and constrained supports.

• Traces and constrained inference — simulate versus generate, Constraint / Address construction, and the trace round-trip between flat position vector and tree-shaped Trace.

• MCMC: MH, HMC, NUTS — Metropolis-Hastings over a Selection, gradient-based chains via mcmc, and log-density-driven chains via mcmc_logpdf.

• Running and resuming chains — run_chain, warmup and checkpointing through MCMCState, save_state / load_state, and posterior summaries with mcmc_summary.