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

Synthetic data drawn from slope = -2, intercept = 10. The xs / ys pair is what we'll condition on.

xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
ys = [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, -9.90]

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]) => [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, …

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}:
 8.23  5.87  3.99  2.59  0.23  -0.66  -3.53  -6.91  -7.24  -9.9

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([:intercept])
  Reactant.ProbProg.Address([:slope])

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#2098)

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}:
  9.41302  -1.91937
  9.56795  -1.88598
  9.98841  -1.96558
 10.5022   -2.04002
  9.64193  -1.91864
 10.8575   -2.04708
 10.5624   -1.97696
  9.94545  -2.02002
  9.5698   -1.84124
 10.1796   -1.98992
  ⋮        
 10.9306   -2.10629
 10.3737   -1.96357
 10.3572   -1.99189
 10.8025   -2.11401
 11.0124   -2.1524
  8.47845  -1.72255
  9.82637  -1.88122
 10.4149   -2.04608
 10.673    -2.09329

Each row is one post-warmup sample; columns hold the sampled values for each selected site:

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

From the sampler output, the posterior mean is:

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

(posterior_mean_intercept = 10.10005867621911, posterior_mean_slope = -1.9717556264200453)

The data were generated from slope = -2, intercept = 10; NUTS recovers both posterior means.

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.13226  10.4134
 -2.09329  11.0906
 -1.87496   9.40189
 -1.96555   9.41682
 -1.99321  10.1236
 -2.02447  10.3259
 -1.99766  10.1869
 -1.99212  10.3043
 -1.92231  10.347
 -1.86255  10.0598
  ⋮        
 -2.13143  10.7526
 -2.04997  10.2702
 -2.07028  10.8019
 -1.98442  10.3042
 -1.94368  10.1219
 -2.03354  10.4207
 -1.88991   9.85419
 -1.94676  10.1356
 -2.02318  10.3356

We get similar inference results

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

(posterior_mean_slope = -1.9848081026726052, posterior_mean_intercept = 10.157414186389726)

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.