Feature request: Stochastic Gradient-MCMC support (SGLD, SGHMC, etc.)

[mattlevine22]

Issue #950 asked about SG-MCMC, and while BlackJAX has kernels (link), using them in NumPyro requires extra glue code for making the log density, transforms, minibatching, and diagnostics (I've been struggling with this a bit, hence my post). Native support would make it much easier to run SG-MCMC directly on NumPyro model---and I think it is a common enough setting that would be worth it. This is valuable not just for minibatching large datasets but also for models with randomized approximate likelihoods (e.g. EnKF, PF; 1901.10568, 1907.06986). Would others be interested in collaborating on adding this?

[JosueT1212]

I would like to collaborate

[deeb01]

I think the missing glue is already there, just undocumented for this case.

plate(subsample_size=B) plus numpyro.subsample already gives an unbiased gradient of the
full-data potential — checked against it over 200 draws at B=100, z-scores 0.5/1.3/0.3.
Selecting the minibatch is a substitute on the plate's index site, and it jits with the
indices traced, so it drops straight into a kernel loop:

def potential(params, indices):
    sub = handlers.substitute(model, {"data": indices})
    return potential_energy(sub, model_args, {}, params)

grad_est = jax.jit(jax.grad(potential))

Transforms and diagnostics are covered by the initialize_model / constrain_fn helpers
from #2203.