What is the best way to measure/benchmark the time it takes to evaluate the log probability of a Numpyro model, as well as taking the gradients with respect to the latent paramters? I am mostly interested in the context of MCMC so I would like to benchmark to the log probability and gradient functions that Numpyro uses when doing inference with
To have a specific example. How would I benchmark the
eight_schools model from the Getting Started section?
import numpy as np import numpyro import numpyro.distributions as dist from jax import random from numpyro.infer import MCMC, NUTS J = 8 y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]) sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]) # Eight Schools example def eight_schools(J, sigma, y=None): mu = numpyro.sample('mu', dist.Normal(0, 5)) tau = numpyro.sample('tau', dist.HalfCauchy(5)) with numpyro.plate('J', J): theta = numpyro.sample('theta', dist.Normal(mu, tau)) numpyro.sample('obs', dist.Normal(theta, sigma), obs=y) nuts_kernel = NUTS(eight_schools) mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000) rng_key = random.PRNGKey(0) mcmc.run(rng_key, J, sigma, y=y, extra_fields=('potential_energy',))