Export a compiled density with jax.export

A compiled FlatPPL function is plain StableHLO inside JAX, so jax.export serializes it, gradient included, into bytes that load and run in a process that has neither FlatPPL nor Enzyme installed. The artefact is the portable model: the same numbers, on any JAX backend, with no compiler in the loop.

The density is the posterior of a Normal likelihood with unknown mean and precision under its conjugate Normal-Gamma prior, so its values, its gradient and its evidence are known in closed form.

Runtime budget: 120 s.

import subprocess
import sys
import tempfile
from math import lgamma, log, pi
from pathlib import Path

import jax
import numpy as np

from flatppl import flatppl

Compile the density

logdensity("posterior") takes mu and log tau as free coordinates.

data = np.array([1.2, 0.4, 2.1, 1.7, 0.9, 1.5, 0.3, 1.1])
model = flatppl(
    """
    tau ~ Gamma(shape = 3.0, rate = 2.0)
    mu ~ Normal(mu = 0.0, sigma = 1.0 / sqrt(2.0 * tau))
    y ~ iid(Normal(mu = mu, sigma = 1.0 / sqrt(tau)), 8)
    observed = external(cartpow(reals, 8))
    prior = lawof(record(mu = mu, tau = tau))
    L = likelihoodof(kernelof(record(y = y), mu = mu, tau = tau), record(y = observed))
    posterior = bayesupdate(L, prior)
    """
).set(observed=data)
density = model.logdensity("posterior")

Serialize the value and the gradient

jax.export traces a jitted function at abstract argument shapes and keeps the lowered program. Serialization needs the flatbuffers package. Here the exported function returns the log density and its gradient on the flat vector of free coordinates.

flat = density.flat
spec = jax.ShapeDtypeStruct(flat.init.shape, flat.init.dtype)
exported = jax.export.export(jax.jit(jax.value_and_grad(flat.fn)))(spec)
artefact = Path(tempfile.mkdtemp()) / "posterior.jax"
artefact.write_bytes(exported.serialize())
print(f"{artefact.name}: {artefact.stat().st_size} bytes, inputs {exported.in_avals}")
print("flat coordinates:", flat.names)

posterior.jax: 6464 bytes, inputs (ShapedArray(float32[2]),)
flat coordinates: ('mu', 'tau')

Reload without FlatPPL

A fresh Python process, isolated from this page’s imports, deserializes the bytes and evaluates them. It asserts that neither flatppl nor enzyme_ad was imported.

points = np.array([[0.3, 0.2], [1.1, -0.4], [0.9, 0.5]], dtype=np.float32)
reload = f"""
import sys
import jax
import numpy as np
exported = jax.export.deserialize(bytearray(open({str(artefact)!r}, "rb").read()))
for point in np.array({points.tolist()!r}, dtype=np.float32):
    value, gradient = exported.call(point)
    print(float(value), *np.asarray(gradient).tolist())
assert not any(name.split(".")[0] in ("flatppl", "enzyme_ad") for name in sys.modules)
"""
output = subprocess.run(
    [sys.executable, "-I", "-c", reload], capture_output=True, text=True, check=True
).stdout
reloaded = np.array([line.split() for line in output.splitlines()], dtype=float)
print(reloaded)

[[-12.75684452   7.57269764  -0.21926591]
 [-12.59423828  -1.20657635   4.44334078]
 [-10.04195213   0.32974482   0.24562703]]

Check against the closed form

The flat vector orders the record fields by name: mu, then tau’s free coordinate u = log tau. With the Jacobian u, the density is log Gamma(tau; 3, 2) + log Normal(mu; 0, 1/sqrt(2 tau)) + sum log Normal(y; mu, 1/sqrt(tau)) + u.

def closed_form(mu: float, u: float) -> float:
    """The posterior log density in free coordinates, written out."""
    tau = np.exp(u)
    prior_tau = 3 * np.log(2) - lgamma(3) + 2 * u - 2 * tau
    prior_mu = 0.5 * np.log(2 * tau / (2 * pi)) - tau * mu**2
    likelihood = 0.5 * data.size * np.log(tau / (2 * pi))
    likelihood -= 0.5 * tau * ((data - mu) ** 2).sum()
    return prior_tau + prior_mu + likelihood + u


def closed_gradient(mu: float, u: float) -> tuple[float, float]:
    """The gradient of :func:`closed_form`."""
    tau = np.exp(u)
    d_mu = -2 * tau * mu + tau * (data - mu).sum()
    d_u = 2 - 2 * tau + 0.5 - tau * mu**2 + 0.5 * data.size + 1
    d_u -= 0.5 * tau * ((data - mu) ** 2).sum()
    return d_mu, d_u


for point, row in zip(points, reloaded, strict=True):
    expected = closed_form(*point)
    assert abs(row[0] - expected) < 1e-5 * abs(expected)
    np.testing.assert_allclose(row[1:], closed_gradient(*point), rtol=1e-4, atol=1e-4)

The artefact keeps the normalization

An exported function has fixed argument shapes, and jax.vmap cannot batch a deserialized one. Export the batched density instead, at the size of a grid over (mu, log tau), and reload it. The grid sum reproduces the closed-form evidence, so no constant was lost on the way out.

a0, b0, k0, m0, n = 3.0, 2.0, 2.0, 0.0, data.size
mean = data.mean()
kn, an = k0 + n, a0 + n / 2
bn = b0 + 0.5 * ((data - mean) ** 2).sum() + k0 * n * (mean - m0) ** 2 / (2 * kn)
log_evidence = (
    lgamma(an) - lgamma(a0) + a0 * log(b0) - an * log(bn)
    + 0.5 * log(k0 / kn) - n / 2 * log(2 * pi)
)  # fmt: skip
mu_axis = np.linspace(-3.0, 5.0, 401, dtype=np.float32)
log_tau_axis = np.linspace(-4.0, 3.0, 401, dtype=np.float32)
grid_points = np.stack(np.meshgrid(mu_axis, log_tau_axis, indexing="ij"), axis=-1)
batch = jax.ShapeDtypeStruct(
    (grid_points.shape[0] * grid_points.shape[1], 2), spec.dtype
)
batched = jax.export.export(jax.jit(jax.vmap(flat.fn)))(batch)
loaded = jax.export.deserialize(bytearray(batched.serialize()))
grid = np.asarray(loaded.call(grid_points.reshape(-1, 2))).reshape(
    grid_points.shape[:2]
)
total = np.trapezoid(np.trapezoid(np.exp(grid), log_tau_axis, axis=1), mu_axis)
print(f"log evidence: artefact {np.log(total):.6f}, closed form {log_evidence:.6f}")
assert abs(np.log(total) - log_evidence) < 1e-6

log evidence: artefact -10.558731, closed form -10.558731