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

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)

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?

Start by running the supplied reproduction script with NUTS and LocScaleReparam across the listed K values. Compare the JAX and NumPyro version combinations to isolate whether the regression belongs to JAX, NumPyro, or their interaction, focusing on compilation and per-step runtime. Done means identifying the cause and documenting a verified workaround or regression boundary.