Package: numpyro
NumPyroProbabilistic programming powered byJAXfor autograd and JIT compilation to GPU/TPU/CPU.Docs and Examples|ForumWhat is NumPyro?NumPyro is a lightweight probabilistic programming library that provides a NumPy backend forPyro. We rely onJAXfor automatic differentiation and JIT compilation to GPU / CPU. NumPyro is under active development, so beware of brittleness, bugs, and changes to the API as the design evolves.NumPyro is designed to belightweightand focuses on providing a flexible substrate that users can build on:Pyro Primitives:NumPyro programs can contain regular Python and NumPy code, in addition toPyro primitiveslikesampleandparam. The model code should look very similar to Pyro except for some minor differences between PyTorch and Numpy's API. See theexamplebelow.Inference algorithms:NumPyro supports a number of inference algorithms, with a particular focus on MCMC algorithms like Hamiltonian Monte Carlo, including an implementation of the No U-Turn Sampler. Additional MCMC algorithms includeMixedHMC(which can accommodate discrete latent variables) as well asHMCECS(which only computes the likelihood for subsets of the data in each iteration). One of the motivations for NumPyro was to speed up Hamiltonian Monte Carlo by JIT compiling the verlet integrator that includes multiple gradient computations. With JAX, we can composejitandgradto compile the entire integration step into an XLA optimized kernel. We also eliminate Python overhead by JIT compiling the entire tree building stage in NUTS (this is possible usingIterative NUTS). There is also a basic Variational Inference implementation together with many flexible (auto)guides for Automatic Differentiation Variational Inference (ADVI). The variational inference implementation supports a number of features, including support for models with discrete latent variables (seeTraceGraph_ELBOandTraceEnum_ELBO).Distributions:Thenumpyro.distributionsmodule provides distribution classes, constraints and bijective transforms. The distribution classes wrap over samplers implemented to work with JAX'sfunctional pseudo-random number generator. The design of the distributions module largely follows fromPyTorch. A major subset of the API is implemented, and it contains most of the common distributions that exist in PyTorch. As a result, Pyro and PyTorch users can rely on the same API and batching semantics as intorch.distributions. In addition to distributions,constraintsandtransformsare very useful when operating on distribution classes with bounded support. Finally, distributions from TensorFlow Probability (TFP) can directly be used in NumPyro models.Effect handlers:Like Pyro, primitives likesampleandparamcan be provided nonstandard interpretations using effect-handlers from thenumpyro.handlersmodule, and these can be easily extended to implement custom inference algorithms and inference utilities.A Simple Example - 8 SchoolsLet us explore NumPyro using a simple example. We will use the eight schools example from Gelman et al., Bayesian Data Analysis: Sec. 5.5, 2003, which studies the effect of coaching on SAT performance in eight schools.The data is given by:>>>importnumpyasnp>>>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]), whereyare the treatment effects andsigmathe standard error. We build a hierarchical model for the study where we assume that the group-level parametersthetafor each school are sampled from a Normal distribution with unknown meanmuand standard deviationtau, while the observed data are in turn generated from a Normal distribution with mean and standard deviation given bytheta(true effect) andsigma, respectively. This allows us to estimate the population-level parametersmuandtauby pooling from all the observations, while still allowing for individual variation amongst the schools using the group-levelthetaparameters.>>>importnumpyro>>>importnumpyro.distributionsasdist>>># Eight Schools example...defeight_schools(J,sigma,y=None):...mu=numpyro.sample('mu',dist.Normal(0,5))...tau=numpyro.sample('tau',dist.HalfCauchy(5))...withnumpyro.plate('J',J):...theta=numpyro.sample('theta',dist.Normal(mu,tau))...numpyro.sample('obs',dist.Normal(theta,sigma),obs=y)Let us infer the values of the unknown parameters in our model by running MCMC using the No-U-Turn Sampler (NUTS). Note the usage of theextra_fieldsargument inMCMC.run. By default, we only collect samples from the target (posterior) distribution when we run inference usingMCMC. However, collecting additional fields like potential energy or the acceptance probability of a sample can be easily achieved by using theextra_fieldsargument. For a list of possible fields that can be collected, see theHMCStateobject. In this example, we will additionally collect thepotential_energyfor each sample.>>>fromjaximportrandom>>>fromnumpyro.inferimportMCMC,NUTS>>>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',))We can print the summary of the MCMC run, and examine if we observed any divergences during inference. Additionally, since we collected the potential energy for each of the samples, we can easily compute the expected log joint density.>>>mcmc.print_summary()# doctest: +SKIPmeanstdmedian5.0%95.0%n_effr_hatmu4.143.183.87-0.769.50115.421.01tau4.123.583.120.518.5690.641.02theta[0]6.406.225.36-2.5415.27176.751.00theta[1]4.965.044.49-1.9814.22217.121.00theta[2]3.655.413.31-3.4713.77247.641.00theta[3]4.475.294.00-3.2212.92213.361.01theta[4]3.224.613.28-3.7210.93242.141.01theta[5]3.894.993.71-3.3912.54206.271.00theta[6]6.555.725.66-1.4315.78124.571.00theta[7]4.815.954.19-3.9013.40299.661.00Numberofdivergences:19>>>pe=mcmc.get_extra_fields()['potential_energy']>>>print('Expected log joint density:{:.2f}'.format(np.mean(-pe)))# doctest: +SKIPExpectedlogjointdensity:-54.55The values above 1 for the split Gelman Rubin diagnostic (r_hat) indicates that the chain has not fully converged. The low value for the effective sample size (n_eff), particularly fortau, and the number of divergent transitions looks problematic. Fortunately, this is a common pathology that can be rectified by using anon-centered paramaterizationfortauin our model. This is straightforward to do in NumPyro by using aTransformedDistributioninstance together with areparameterizationeffect handler. Let us rewrite the same model but instead of samplingthetafrom aNormal(mu, tau), we will instead sample it from a baseNormal(0, 1)distribution that is transformed using anAffineTransform. Note that by doing so, NumPyro runs HMC by generating samplestheta_basefor the baseNormal(0, 1)distribution instead. We see that the resulting chain does not suffer from the same pathology — the Gelman Rubin diagnostic is 1 for all the parameters and the effective sample size looks quite good!>>>fromnumpyro.infer.reparamimportTransformReparam>>># Eight Schools example - Non-centered Reparametrization...defeight_schools_noncentered(J,sigma,y=None):...mu=numpyro.sample('mu',dist.Normal(0,5))...tau=numpyro.sample('tau',dist.HalfCauchy(5))...withnumpyro.plate('J',J):...withnumpyro.handlers.reparam(config={'theta':TransformReparam()}):...theta=numpyro.sample(...'theta',...dist.TransformedDistribution(dist.Normal(0.,1.),...dist.transforms.AffineTransform(mu,tau)))...numpyro.sample('obs',dist.Normal(theta,sigma),obs=y)>>>nuts_kernel=NUTS(eight_schools_noncentered)>>>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',))>>>mcmc.print_summary(exclude_deterministic=False)# doctest: +SKIPmeanstdmedian5.0%95.0%n_effr_hatmu4.083.514.14-1.699.71720.431.00tau3.963.313.090.018.34488.631.00theta[0]6.485.726.08-2.5314.96801.591.00theta[1]4.955.104.91-3.7012.821183.061.00theta[2]3.655.583.72-5.7112.13581.311.00theta[3]4.565.044.32-3.1412.921282.601.00theta[4]3.414.793.47-4.1610.79801.251.00theta[5]3.584.803.78-3.9511.551101.331.00theta[6]6.315.175.75-2.9313.871081.111.00theta[7]4.815.384.61-3.2914.05954.141.00theta_base[0]0.410.950.40-1.091.95851.451.00theta_base[1]0.150.950.20-1.421.661568.111.00theta_base[2]-0.080.98-0.10-1.681.541037.161.00theta_base[3]0.060.890.05-1.421.471745.021.00theta_base[4]-0.140.94-0.16-1.651.45719.851.00theta_base[5]-0.100.96-0.14-1.571.511128.451.00theta_base[6]0.380.950.42-1.321.821026.501.00theta_base[7]0.100.970.10-1.511.651190.981.00Numberofdivergences:0>>>pe=mcmc.get_extra_fields()['potential_energy']>>># Compare with the earlier value>>>print('Expected log joint density:{:.2f}'.format(np.mean(-pe)))# doctest: +SKIPExpectedlogjointdensity:-46.09Note that for the class of distributions withloc,scaleparameters such asNormal,Cauchy,StudentT, we also provide aLocScaleReparamreparameterizer to achieve the same purpose. The corresponding code will bewith numpyro.handlers.reparam(config={'theta': LocScaleReparam(centered=0)}):
    theta = numpyro.sample('theta', dist.Normal(mu, tau))Now, let us assume that we have a new school for which we have not observed any test scores, but we would like to generate predictions. NumPyro provides aPredictiveclass for such a purpose. Note that in the absence of any observed data, we simply use the population-level parameters to generate predictions. ThePredictiveutility conditions the unobservedmuandtausites to values drawn from the posterior distribution from our last MCMC run, and runs the model forward to generate predictions.>>>fromnumpyro.inferimportPredictive>>># New School...defnew_school():...mu=numpyro.sample('mu',dist.Normal(0,5))...tau=numpyro.sample('tau',dist.HalfCauchy(5))...returnnumpyro.sample('obs',dist.Normal(mu,tau))>>>predictive=Predictive(new_school,mcmc.get_samples())>>>samples_predictive=predictive(random.PRNGKey(1))>>>print(np.mean(samples_predictive['obs']))# doctest: +SKIP3.9886456More ExamplesFor some more examples on specifying models and doing inference in NumPyro:Bayesian Regression in NumPyro- Start here to get acquainted with writing a simple model in NumPyro, MCMC inference API, effect handlers and writing custom inference utilities.Time Series Forecasting- Illustrates how to convert for loops in the model to JAX'slax.scanprimitive for fast inference.Annotation examples- Illustrates how to utilize the enumeration mechanism to perform inference for models with discrete latent variables.Baseball example- Using NUTS for a simple hierarchical model. Compare this with the baseball example inPyro.Hidden Markov Modelin NumPyro as compared toStan.Variational Autoencoder- As a simple example that uses Variational Inference with neural networks.Pyro implementationfor comparison.Gaussian Process- Provides a simple example to use NUTS to sample from the posterior over the hyper-parameters of a Gaussian Process.Horseshoe Regression- Shows how to implemement generalized linear models equipped with a Horseshoe prior for both binary-valued and real-valued outputs.Statistical Rethinking with NumPyro-Notebookscontaining translation of the code in Richard McElreath'sStatistical Rethinkingbook second version, to NumPyro.Other model examples can be found in theexamplessite.Pyro users will note that the API for model specification and inference is largely the same as Pyro, including the distributions API, by design. However, there are some important core differences (reflected in the internals) that users should be aware of. e.g. in NumPyro, there is no global parameter store or random state, to make it possible for us to leverage JAX's JIT compilation. Also, users may need to write their models in a morefunctionalstyle that works better with JAX. Refer toFAQsfor a list of differences.Overview of inference algorithmsWe provide an overview of most of the inference algorithms supported by NumPyro and offer some guidelines about which inference algorithms may be appropriate for different classes of models.MCMCNUTS, which is an adaptive variant ofHMC, is probably the most commonly used inference algorithm in NumPyro. Note that NUTS and HMC are not directly applicable to models with discrete latent variables, but in cases where the discrete variables have finite support and summing them out (i.e. enumeration) is tractable, NumPyro will automatically sum out discrete latent variables and perform NUTS/HMC on the remaining continuous latent variables.
As discussed above, modelreparameterizationmay be important in some cases to get good performance. Note that, generally speaking, we expect inference to be harder as the dimension of the latent space increases. See thebad geometrytutorial for additional tips and tricks.MixedHMCcan be an effective inference strategy for models that contain both continuous and discrete latent variables.HMCECScan be an effective inference strategy for models with a large number of data points. It is applicable to models with continuous latent variables. Seeherefor an example.BarkerMHis a gradient-based MCMC method that may be competitive with HMC and NUTS for some models. It is applicable to models with continuous latent variables.HMCGibbscombines HMC/NUTS steps with custom Gibbs updates. Gibbs updates must be specified by the user.DiscreteHMCGibbscombines HMC/NUTS steps with Gibbs updates for discrete latent variables. The corresponding Gibbs updates are computed automatically.SAis the only MCMC method in NumPyro that does not leverage gradients. It is only applicable to models with continuous latent variables. It is expected to perform best for models whose latent dimension is low to moderate. It may be a good choice for models with non-differentiable log densities. Note that SA generally requires averylarge number of samples, as mixing tends to be slow. On the plus side individual steps can be fast.Like HMC/NUTS, all remaining MCMC algorithms support enumeration over discrete latent variables if possible (seerestrictions). Enumerated sites need to be marked withinfer={'enumerate': 'parallel'}like in theannotation example.Nested SamplingNestedSampleroffers a wrapper forjaxns. SeeJAXNS's readthedocsfor examples andNested Sampling for Gaussian Shellsexample for how to apply the sampler on numpyro models. Can handle arbitrary models, including ones with discrete RVs, and non-invertible transformations.Stochastic variational inferenceVariational objectivesTrace_ELBOis our basic ELBO implementation.TraceMeanField_ELBOis likeTrace_ELBObut computes part of the ELBO analytically if doing so is possible.TraceGraph_ELBOoffers variance reduction strategies for models with discrete latent variables. Generally speaking, this ELBO should always be used for models with discrete latent variables.TraceEnum_ELBOoffers variable enumeration strategies for models with discrete latent variables. Generally speaking, this ELBO should always be used for models with discrete latent variables when enumeration is possible.Automatic guides (appropriate for models with continuous latent variables)AutoNormalandAutoDiagonalNormalare our basic mean-field guides. If the latent space is non-euclidean (due to e.g. a positivity constraint on one of the sample sites) an appropriate bijective transformation is automatically used under the hood to map between the unconstrained space (where the Normal variational distribution is defined) to the corresponding constrained space (note this is true for all automatic guides). These guides are a great place to start when trying to get variational inference to work on a model you are developing.AutoMultivariateNormalandAutoLowRankMultivariateNormalalso construct Normal variational distributions but offer more flexibility, as they can capture correlations in the posterior. Note that these guides may be difficult to fit in the high-dimensional setting.AutoDeltais used for computing point estimates via MAP (maximum a posteriori estimation). Seeherefor example usage.AutoBNAFNormalandAutoIAFNormaloffer flexible variational distributions parameterized by normalizing flows.AutoDAISis a powerful variational inference algorithm that leverages HMC. It can be a good choice for dealing with highly correlated posteriors but may be computationally expensive depending on the nature of the model.AutoSurrogateLikelihoodDAISis a powerful variational inference algorithm that leverages HMC and that supports data subsampling.AutoSemiDAISconstructs a posterior approximation likeAutoDAISfor local latent variables but provides support for data subsampling during ELBO training by utilizing a parametric guide for global latent variables.AutoLaplaceApproximationcan be used to compute a Laplace approximation.Stein Variational InferenceSee thedocsfor more details.InstallationLimited Windows Support:Note that NumPyro is untested on Windows, and might require building jaxlib from source. See thisJAX issuefor more details. Alternatively, you can installWindows Subsystem for Linuxand use NumPyro on it as on a Linux system. See alsoCUDA on Windows Subsystem for Linuxandthis forum postif you want to use GPUs on Windows.To install NumPyro with the latest CPU version of JAX, you can use pip:pip install numpyroIn case of compatibility issues arise during execution of the above command, you can instead force the installation of a known
compatible CPU version of JAX withpip install numpyro[cpu]To useNumPyro on the GPU, you need to install CUDA first and then use the following pip command:pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlIf you need further guidance, please have a look at theJAX GPU installation instructions.To runNumPyro on Cloud TPUs, you can look at someJAX on Cloud TPU examples.For Cloud TPU VM, you need to setup the TPU backend as detailed in theCloud TPU VM JAX Quickstart Guide.
After you have verified that the TPU backend is properly set up,
you can install NumPyro using thepip install numpyrocommand.Default Platform:JAX will use GPU by default if CUDA-supportedjaxlibpackage is installed. You can useset_platformutilitynumpyro.set_platform("cpu")to switch to CPU at the beginning of your program.You can also install NumPyro from source:git clone https://github.com/pyro-ppl/numpyro.git
cd numpyro
# install jax/jaxlib first for CUDA support
pip install -e .[dev]  # contains additional dependencies for NumPyro developmentYou can also install NumPyro with conda:conda install -c conda-forge numpyroFrequently Asked QuestionsUnlike in Pyro,numpyro.sample('x', dist.Normal(0, 1))does not work. Why?You are most likely using anumpyro.samplestatement outside an inference context. JAX does not have a global random state, and as such, distribution samplers need an explicit random number generator key (PRNGKey) to generate samples from. NumPyro's inference algorithms use theseedhandler to thread in a random number generator key, behind the scenes.Your options are:Call the distribution directly and provide aPRNGKey, e.g.dist.Normal(0, 1).sample(PRNGKey(0))Provide therng_keyargument tonumpyro.sample. e.g.numpyro.sample('x', dist.Normal(0, 1), rng_key=PRNGKey(0)).Wrap the code in aseedhandler, used either as a context manager or as a function that wraps over the original callable. e.g.withhandlers.seed(rng_seed=0):# random.PRNGKey(0) is usedx=numpyro.sample('x',dist.Beta(1,1))# uses a PRNGKey split from random.PRNGKey(0)y=numpyro.sample('y',dist.Bernoulli(x))# uses different PRNGKey split from the last one, or as a higher order function:deffn():x=numpyro.sample('x',dist.Beta(1,1))y=numpyro.sample('y',dist.Bernoulli(x))returnyprint(handlers.seed(fn,rng_seed=0)())Can I use the same Pyro model for doing inference in NumPyro?As you may have noticed from the examples, NumPyro supports all Pyro primitives likesample,param,plateandmodule, and effect handlers. Additionally, we have ensured that thedistributionsAPI is based ontorch.distributions, and the inference classes likeSVIandMCMChave the same interface. This along with the similarity in the API for NumPy and PyTorch operations ensures that models containing Pyro primitive statements can be used with either backend with some minor changes. Example of some differences along with the changes needed, are noted below:Anytorchoperation in your model will need to be written in terms of the correspondingjax.numpyoperation. Additionally, not alltorchoperations have anumpycounterpart (and vice-versa), and sometimes there are minor differences in the API.pyro.samplestatements outside an inference context will need to be wrapped in aseedhandler, as mentioned above.There is no global parameter store, and as such usingnumpyro.paramoutside an inference context will have no effect. To retrieve the optimized parameter values from SVI, use theSVI.get_paramsmethod. Note that you can still useparamstatements inside a model and NumPyro will use thesubstituteeffect handler internally to substitute values from the optimizer when running the model in SVI.PyTorch neural network modules will need to rewritten asstaxneural networks. See theVAEexample for differences in syntax between the two backends.JAX works best with functional code, particularly if we would like to leverage JIT compilation, which NumPyro does internally for many inference subroutines. As such, if your model has side-effects that are not visible to the JAX tracer, it may need to rewritten in a more functional style.For most small models, changes required to run inference in NumPyro should be minor. Additionally, we are working onpyro-apiwhich allows you to write the same code and dispatch it to multiple backends, including NumPyro. This will necessarily be more restrictive, but has the advantage of being backend agnostic. See thedocumentationfor an example, and let us know your feedback.How can I contribute to the project?Thanks for your interest in the project! You can take a look at beginner friendly issues that are marked with thegood first issuetag on Github. Also, please feel to reach out to us on theforum.Future / Ongoing WorkIn the near term, we plan to work on the following. Please open new issues for feature requests and enhancements:Improving robustness of inference on different models, profiling and performance tuning.Supporting more functionality as part of thepyro-apigeneric modeling interface.More inference algorithms, particularly those that require second order derivatives or use HMC.Integration withFunsorto support inference algorithms with delayed sampling.Other areas motivated by Pyro's research goals and application focus, and interest from the community.Citing NumPyroThe motivating ideas behind NumPyro and a description of Iterative NUTS can be found in thispaperthat appeared in NeurIPS 2019 Program Transformations for Machine Learning Workshop.If you use NumPyro, please consider citing:@article{phan2019composable,
  title={Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro},
  author={Phan, Du and Pradhan, Neeraj and Jankowiak, Martin},
  journal={arXiv preprint arXiv:1912.11554},
  year={2019}
}as well as@article{bingham2019pyro,
  author    = {Eli Bingham and
               Jonathan P. Chen and
               Martin Jankowiak and
               Fritz Obermeyer and
               Neeraj Pradhan and
               Theofanis Karaletsos and
               Rohit Singh and
               Paul A. Szerlip and
               Paul Horsfall and
               Noah D. Goodman},
  title     = {Pyro: Deep Universal Probabilistic Programming},
  journal   = {J. Mach. Learn. Res.},
  volume    = {20},
  pages     = {28:1--28:6},
  year      = {2019},
  url       = {http://jmlr.org/papers/v20/18-403.html}
}