NUTS ~100–200× slower on numpyro 0.21.0 than 0.20.0 for a simple hierarchical Bernoulli model (GPU)

[yingboli]

Bug Description

After upgrading from numpyro 0.20.0 → 0.21.0, NUTS becomes dramatically slower
for a simple Bayesian hierarchical model with a binary (Bernoulli) likelihood, when
the number of groups K is small relative to the number of observations N.

The slowdown is not caused by the sampler taking more leapfrog steps: the mean
number of steps per sample is essentially unchanged across versions. Only the
wall-clock runtime explodes, which points at a regression in per-step / compilation
cost rather than sampler behavior.

The upgrade also moves jax from 0.9.1 → 0.10.2, so the root cause could live in
either package.

Results (single L4 GPU, N = 100000, 1000 warmup + 2000 samples, 1 chain)

jax	numpyro	K	runtime_s	mean_num_steps
0.9.1	0.20.0	8	26.76	23.66
0.9.1	0.20.0	48	23.70	28.76
0.9.1	0.20.0	900	18.19	15
0.9.1	0.20.0	15000	38.16	31
0.10.2	0.21.0	8	5101.77	25.86
0.10.2	0.21.0	48	1554.13	29.15
0.10.2	0.21.0	900	40.66	15
0.10.2	0.21.0	15000	40.15	31

Observations:

• At K=8 runtime goes from 27s to 5100s (190×); at K=48, 24s → 1550s (65×).
• mean_num_steps is basically identical between versions (e.g. 23.66 vs 25.86 at
  K=8), so the sampler is doing the same amount of work — each step is just far
  slower.
• The regression shrinks as K grows: at K=900 and K=15000 the two versions are
  comparable. The blow-up is worst when K is small (few groups, many observations).

Steps to Reproduce

import time
import csv

import numpy as np
import jax
import jax.numpy as jnp

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
from numpyro.infer.reparam import LocScaleReparam

N = 100000

def bernoulli_hier(group_idx, K, y=None):
    tausq = numpyro.sample("tausq", dist.InverseGamma(5.0, 0.5))
    with numpyro.plate("levels", K):
        with numpyro.handlers.reparam(config={"alpha": LocScaleReparam(centered=0)}):
            alpha = numpyro.sample("alpha", dist.Normal(0.0, jnp.sqrt(tausq)))
    numpyro.sample("obs", dist.Bernoulli(logits=alpha[group_idx]), obs=y)

def make_bernoulli_data(N, K, seed=42):
    rng = np.random.default_rng(seed)
    group_idx = rng.integers(0, K, size=N).astype(np.int32)
    alpha_true = rng.normal(0.0, 0.5, size=K)
    p = 1.0 / (1.0 + np.exp(-alpha_true[group_idx]))
    y = (rng.random(size=N) < p).astype(np.int32)
    return jnp.asarray(group_idx), jnp.asarray(y)

def run_fit(model_fn, **model_kwargs):
    mcmc = MCMC(NUTS(model_fn), num_warmup=1000, num_samples=2000, num_chains=1)
    t0 = time.time()
    mcmc.run(jax.random.PRNGKey(42), extra_fields=("num_steps",), **model_kwargs)
    runtime = time.time() - t0
    ns = np.asarray(mcmc.get_extra_fields()["num_steps"])
    return runtime, ns

rows = []
for K in [8, 48, 900, 15000]:
    group_idx, y = make_bernoulli_data(N, K)
    runtime, ns = run_fit(bernoulli_hier, group_idx=group_idx, K=K, y=y)
    row = {
        "jax_version": jax.__version__,
        "numpyro_version": numpyro.__version__,
        "K": K,
        "runtime_s": round(runtime, 2),
        "mean_num_steps": round(float(ns.mean()), 2),
    }
    rows.append(row)
    print(row)

fieldnames = ["jax_version", "numpyro_version", "K", "runtime_s", "mean_num_steps"]
with open("numpyro_issue_results_local.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)

Environment

• GPU: NVIDIA L4
• numpyro: 0.20.0 vs 0.21.0
• jax / jaxlib: 0.9.1 (with 0.20.0) vs 0.10.2 (with 0.21.0)
• Python: 3.12.13
• CUDA / cuDNN / driver: cuda 13
• OS: Ubuntu 22.04.5 LTS

Expected Behavior

• Is this a known regression introduced by the 0.21.0 changes, the jax 0.10.x
  upgrade, or the interaction of the two?
• Is there a workaround (e.g. a config flag, a different reparameterization, or
  pinning a specific jax version) to recover 0.20.0 performance on 0.21.0?

[JosueT1212]

Can I work on this?

[Qazalbash]

Can I work on this?

Yes, please! If possible, we can add this to the benchmarks too.