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?
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?