Inspecting automatic transformations in MCMC

numpyro
1

I’m a numpyro newbie so would like some help confirming the details of automatic transformations in MCMC.

I want to use HMC in numpyro for inferring posterior parameter distributions using constrained uniform priors. I understand from the docs numpyro will automatically transform the priors from constrained to unconstrained and then automatically do the inverse transform when I request the samples. I expect the transform to be a composition of an affine and sigmoid transformation. How can I verify this in the code?

Here is my mwe:

import jax.numpy as jnp
import jax.random as jnpr
from numpyro import (distributions as dist, infer, sample, set_host_device_count)
from numpyro.infer.util import get_transforms

def make_synthetic_data(rng, true_values, obs_sigma_sq, N):
    lower, upper = 0, 10
    x = np.sort(rng.uniform(lower, upper, N))
    y = true_values["m1"] * x + true_values["c"]
    yerr = np.full(N, obs_sigma_sq)
    y += yerr * rng.standard_normal(size=N)
    return x, y, yerr

def linear_model(X, yerr, y=None):
    c = sample("c", dist.Uniform(-1.0, 1.0))
    m1 = sample("m1", dist.Uniform(0.1, 5.0))
    theta = jnp.array([c, m1])
    mu = jnp.dot(X, theta)
    sample("obs", dist.Normal(mu, yerr), obs=y)

seed = 1234
rng = np.random.default_rng(seed)
key = jnpr.PRNGKey(seed)
set_host_device_count(2)
N = 50
unknowns = {"c": (-1, 1), "m1": (0.1, 5.0)}
tvs = {"c": 0.2, "m1": 2.1}

x, y_obs, yerr = make_synthetic_data(rng, tvs, 0.01, N)
X_obs = np.vander(x, 2, increasing=True)
sampler = infer.MCMC(
    infer.NUTS(linear_model),
   num_warmup=2000,
   num_samples=2000,
   num_chains=2,
   progress_bar=True,
)
sampler.run(key, X_obs, yerr, y=y_obs)
samples = sampler.get_samples()["c"]

transforms = get_transforms(model=linear_model, model_args=(X_obs, yerr, y_obs),
                            model_kwargs={}, params=tvs)

x_unconstrained = 100.0
x_constrained = transforms["c"](x_unconstrained)
print(f"{x_unconstrained} -> {x_constrained}")

The code shows how to retrieve samples from the constrained space and how to retrieve the composed transformations for each parameter. My specific questions are:

  • How to get samples in unconstrained space? Do I have to apply the inverse transform or can I retrieve them directly from sampler?
  • transforms["c"] is a composition. How can I decompose this so I can verify an affine and sigmoid transformation have been used?
  • Where in the numpyro code are the transforms identified and composed?

Any help would be greatly appreciated.