PyMC Comparison

This page maps PyMC and BlackJAX concepts to their Norn equivalents. Norn's design is closest to BlackJAX: both provide functional kernel APIs where the sampler state is explicit and the log-density function is passed at each step.

One-Line Sampling

PyMC:

import pymc as pm

with pm.Model():
    x = pm.Normal("x", mu=0, sigma=1, shape=2)
    trace = pm.sample(1000, tune=500)

Norn:

let log_prob x = Nx.mul_s (Nx.sum (Nx.square x)) (-0.5) in
let init = Nx.zeros Nx.float64 [| 2 |] in
let result = Norn.nuts ~n:1000 ~num_warmup:500 log_prob init

PyMC builds a probabilistic model and derives the log-density automatically. Norn takes the log-density function directly -- you write it yourself or build it from your model. Rune handles the gradient.

BlackJAX Kernel API

BlackJAX:

import blackjax
import jax

kernel = blackjax.nuts(log_prob, step_size=0.5)
state = kernel.init(jax.numpy.zeros(2))

for _ in range(1000):
    key, subkey = jax.random.split(key)
    state, info = kernel.step(subkey, state)

Norn:

let metric = Norn.unit_metric 2 in
let kernel = Norn.nuts_kernel ~step_size:0.5 ~metric () in
let state = ref (kernel.init (Nx.zeros Nx.float64 [| 2 |]) log_prob) in

for _ = 1 to 1000 do
  let new_state, _info = kernel.step !state log_prob in
  state := new_state
done

Both use a {init; step} pattern. The key difference: BlackJAX threads a PRNG key explicitly, while Norn uses Nx's RNG context (Rng.run).

Adaptation

BlackJAX:

warmup = blackjax.window_adaptation(blackjax.nuts, log_prob)
state, kernel, _ = warmup.run(key, jax.numpy.zeros(2), 1000)

Norn:

(* Adaptation is built into sample/nuts/hmc *)
let result = Norn.nuts ~n:1000 ~num_warmup:500 log_prob init

(* Or use sample for control over the kernel *)
let result =
  Norn.sample ~n:1000 ~num_warmup:500 log_prob init (fun ~step_size ~metric ->
      Norn.nuts_kernel ~step_size ~metric ())

In BlackJAX, adaptation is a separate step that returns a tuned kernel. In Norn, adaptation is integrated into sample -- it adapts step size and mass matrix during warmup, then freezes them for sampling.

Samplers

Integrators

Usage comparison:

# BlackJAX
kernel = blackjax.nuts(log_prob, step_size=0.5,
                       integrator=blackjax.mcmc.integrators.mclachlan)

(* Norn *)
let kernel =
  Norn.nuts_kernel ~integrator:Norn.mclachlan ~step_size:0.5 ~metric ()

Metrics (Mass Matrix)

Diagnostics

State and Info

BlackJAX state:

state.position      # current sample
state.logdensity    # log p(x)
state.logdensity_grad  # grad log p(x)

Norn state:

state.position         (* Nx.float64_t, shape [dim] *)
state.log_density      (* float *)
state.grad_log_density (* Nx.float64_t, shape [dim] *)

BlackJAX info:

info.acceptance_rate
info.is_divergent
info.energy
info.num_integration_steps

Norn info:

info.acceptance_rate        (* float in [0, 1] *)
info.is_divergent           (* bool *)
info.energy                 (* float *)
info.num_integration_steps  (* int *)

Key Differences