0.3.0

Breaking Changes

• HMC's find_heuristic_step_size (this functionality is different from step size adaptation scheme) is disabled by default to improve compiling time. Previous behavior can be enabled by setting find_heuristic_step_size=True.
• The automatic reparameterization mechanism introduced in NumPyro 0.2 is removed, in favor of reparam handler. See the eight schools example for the new usage pattern.
• Automatic Guide Generation module is moved from numpyro.contrib.autoguide to the main inference module numpyro.infer.autoguide.
• Various API changes to match Pyro API:
  • mask handler: mask_array arg is renamed to mask.
  • scale handler: scale_factor arg is renamed to scale.
  • condition and substitute handlers: param_map is renamed to data.
  • MultivariateAffineTransform transform is renamed to LowerCholeskyAffine.
  • init_to_prior strategy is renamed to init_to_sample.

New Features

• Funsor-based NumPyro: allow enumeration over discrete latent variables. See mixture and Markov examples below for some applications.
• New primitives: plate_stack and scan. If your model has Python for loop, considering using scan instead to improve compiling time.
• New handlers: reparam, scope, and lift.
• New distributions: von Mises, Gumbel, Logistic, Laplace, TruncatedPolyaGamma, ExpandedDistribution, MaskedDistribution, and ImproperUniform.
• Distribution has new properties is_discrete, has_enumerate_support, and new methods shape, enumerate_support, expand, expand_by, mask. In addition, Distribution has been registered as a JAX Pytree class, with corresponding methods tree_flatten and tree_unflatten.
• New constraint: less_than.
• Port Tensor Indexing from Pyro.
• Port some Reparameterizers from Pyro.
• Add batch_ndims arg to Predictive and log_likelihood to allow using those utilities with arbitrary number of batch dimensions.

New Examples

• Proportion Test example: You are managing a business and want to test if calling your customers will increase their chance of making a purchase.
• Bayesian Models of Annotation examples: illustrates enumeration for mixture models.
• Enumerate HMM examples: illustrates enumeration for markov models.
• Bayesian Imputation tutorial.

Enhancements and Bug Fixes

• HMC/NUTS compiling time is greatly improved, especially for large models.
• More efficient BTRS algorithm for sampling from Binomial distribution. #537
• Allow arbitrary order of plate statements. #555
• Fix KeyError with scale handler and deterministic primitive. #577
• Fix Poisson sampler entering into infinite loop under vmap. #582
• Fix the double compilation issue in numpyro.optim classes. #603
• Use ExpandedDistribution in numpyro.plate. #616
• Timeseries forecasting tutorial is updated with scan primitive and the usage of Predictive for forecasting. #608 #657
• Tweak sparse regression example to bring the model into exact alignment with the reference. #669
• Add MetropolisHastings algorithm as an example of MCMCKernel. #680

Thanks Nikolaos @daydreamt, Daniel Sheldon @dsheldon, Lukas Prediger @lumip, Freddy Boulton @freddyaboulton, Wouter van Amsterdam @vanAmsterdam, and many others for their contributions and helpful feedback!