0.7.1

In 0.7.0 release, the wheel file uploaded to PyPI had some files not updated. This release fixes that issue.

0.7.0

Since this release, NumPyro can be installed along with the latest jax and jaxlib releases (their version restrictions have been relaxed). In addition, NumPyro will use the default JAX platform so if you installed JAX with GPU/TPU support, their devices will be used by default.

New Features

• New distributions: SoftLaplace, Weibull, BetaProportion, NegativeBinomial, NegativeBinomial2, ZeroInflatedDistribution, ZeroInflatedPoisson, ZeroInflatedNegativeBinomial2, FoldedDistribution
• Support for DeepMind's Optax optimizers in SVI
• New ELBO objective: TraceGraph_ELBO for non-reparameterized latent variables (e.g. discrete latent variables)
• A new wrapper NestedSampler to leverage the nested sampling package jaxns for NumPyro models
• Implement cdf and icdf methods for many distributions
• New cond primitive.
• New infer_discrete handler to sample discrete sites under enumeration. Check out the annotation example for a usage.
• Structural mass matrix can be specified via dense_mass argument of the HMC/NUTS constructor #963
• New examples:
  • Thompson sampling for Bayesian optimization with GPs
  • Latent Dirichlet Allocation for topic modeling: a great example to illustrate the usage of Flax/Haiku in NumPyro

Enhancements and Bug Fixes

• Documentation and examples are greatly enhanced to make features more accessible
• Fix chain detection for various CPU device strings #1077
• Fix AutoNormal's quantiles method for models with non-scalar latent sites #1066
• Fix LocScaleReparam with center=1 #1059
• Enhance auto guides to support models with deterministic sites #1022
• Support for mutable states in Flax and Haiku modules #1016
• Fix a bug in auto guides that happens when using the guide in Predictive #1013
• Support decorator syntax for effect handlers #1009
• Implement sparse Poisson log probability #1003
• Support total_count=0 in Multinomial distribution #1000
• Add a flag to control regularize mass matrix behavior in mass matrix adaptation #998
• Add experimental Dockerfiles #996
• Allow setting max tree depth of NUTS sampler during warmup phase #984
• Fix dimensions mixed up in ExpandedDistribution.sample method #972
• MCMC objects can be pickled now #968

This release is made of great contributions and feedbacks from the Pyro community: @ahoho, @kpj, @gustavehug, @AndrewCSQ, @jatentaki, @tcbegley, @dominikstrb, @justinrporter, @dirmeier, @irustandi, @MarcoGorelli, @lumip, and many others. Thank you!

0.6.0

New Features

• Progress bar is available for running parallel MCMC chains.
• New samplers:
  • BarkerMH - a Metropolis-Hastings sampler that uses a skew-symmetric proposal distribution that depends on the gradient of the potential
  • New taylor_proxy for HMCECS sampler. This control variate significantly improves the performance of HMCECS on tall data.
  • MixedHMC for mixed discrete and continuous variables
• New distributions:
  • ProjectedNormal is similar to von Mises and von Mises-Fisher distributions but permits tractable variational inference via reparametrizers
  • TruncatedDistribution to truncate over a family of symmetric distributions: Cauchy, Laplace, Logistic, Normal, StudentT
• New method Distribution.infer_shapes() for static shape analysis.
• New constraints: sphere, positive_ordered_vector, softplus_positive, softplus_lower_cholesky
• New transforms: SoftplusTransform, SoftplusLowerCholeskyTransform
• New reparameterizer: ProjectedNormalReparam for ProjectedNormal distribution
• New obs_mask argument in sample primitive for masked conditioning
• New examples:
  • Hamiltonian Monte Carlo with Energy Conserving Subsampling
  • MCMC Methods for Tall Data

Enhancements and Bug Fixes

• Improve precision for Dirichlet distributions with small concentration #943
• Make it easy to use softplus transforms in autoguides #941
• Improving compiling time in MCMC samplers - compiling time is 2x faster than previously #924
• Reduce memory requirement for AutoLowRankMultivariateNormal.quantiles #921
• Example of how to use Distribution.mask #917
• Add goodness of fit helpers for testing distributions #916
• Enabling sampling with intermediates for ExpandedDistribution #909
• Fix DiscreteHMCGibbs to work with multiple chains #908
• Fix missing infer key in handlers.lift #892

Thanks @loopylangur, Dominik Straub @dominikstrb, Jeremie Coullon @jeremiecoullon, Ola Rønning @OlaRonning, Lukas Prediger @lumip, Raúl Peralta Lozada @RaulPL, Vitalii Kleshchevnikov @vitkl, Matt Ludkin @ludkinm, and many others for your contributions and feedback!

0.5.0

New documentation page with galleries of tutorials and examples num.pyro.ai.

New Features

• New primitive: prng_key to draw a random key under seed handler.
• New autoguide: AutoDelta
• New samplers:
  • HMCGibbs: a general HMC/NUTS-within-Gibbs interface.
  • DiscreteHMCGibbs: HMC/NUTS-within-Gibbs for models with discrete latent variables.
  • HMCECS: HMC/NUTS with energy conserving subsampling.
• New example:
  • CJS capture-recapture model for ecological data
• New kernels module in numpyro.contrib.einstein, in preparing for (Ein)Stein VI inference in future releases.
• New user-friendly SVI.run method to simplify the training phase of SVI inference.
• New feasible_like method in constraints.
• New methods forward_shape and inverse_shape in Transform to infer output shape given input shape.
• Transform.inv now returns an inversed transform, hence enables many new (inversed) transforms.
• Support thinning in MCMC.
• Add post_warmup_state and last_state to allow sequential sampling strategy in MCMC: allow to keep calling .run method to get more samples.
• New history argument to support for Markov models with history > 1 in scan.
• New forward_model_differentiation argument in HMC/NUTS kernels to allow to use forward mode differentiation.

Enhancements and Bug Fixes

• #886 Make TransformReparam compatible with .to_event()
• #883 Improve gradient computation of Euclidean kinetic energy.
• #872 Enhance masked distribution to allow gradient propagate properly when using mask handler for invalid data.
• #865 Make subsample faster in CPU.
• #860 Fix for memory leak in MCMC.
• #849 Expose logits attribute to some discrete distributions
• #848 Add has_rsample and rsample attribute to distributions
• #832 Allow a callable to return an init value in param primitive
• #824 Fix for cannot using sample method of TFP distributions in sample primitive.
• #823 Demo on how to use various init strategies in Gaussian Process example.
• #822 Allow haiku/flax modules to take general args/kwargs in init.
• #821 Better error messages when rng_key is missing.
• #818 Better error messages when an error happens in the middle of inference.
• #805 Display correct progress bar message after running MCMC.warmup.
• #801 Raise an error early if missing plates for models with discrete latent variables.
• #797 MCMC vectorized chain method works for models with deterministic sites.
• #796 Bernoulli distribution returns an int instead of a boolean.
• #795 Reveal signature for help(Distribution).

Thanks Ola Ronning @OlaRonning, Armin Stepanjan @ab-10, @cerbelaut, Xi Wang @xidulu, Wouter van Amsterdam @vanAmsterdam, @loopylangur, and many others for your contributions and helpful feedback!

0.4.1

New Features

• #772 Add DirichletMultinomial distribution.
• #773 (experiment) Add collapse handler to exploit conjugacy relations.

Enhancements and Bug Fixes

• #764 Make exception chaining more user-friendly. Thanks, @akihironitta!
• #766 Relax interval constraint.
• #776 Fix bugs in methods log_prob and sample of VonMises distribution.
• #775 Make validation mechanism compatible with omnistaging since JAX 0.2.
• #780 Fix name dimensions of sample sites under contrib.funsor's plate handler.

0.4.0

Experimental integrations with JAX-based TensorFlow Probability and neural network libraries Flax and Haiku. New high-quality tutorials written by NumPyro contributors. JAX 0.2 enables "omnistaging" by default (see this guide for what omnistaging means and how to update your code if it is broken after the upgrade - you can also disable this new behavior with jax.config.disable_omnistaging()).

New Features

• New primitives in numpyro.contrib.module to create Bayesian Neural Networks (BNN) using Flax or Haiku: flax_module, random_flax_module, haiku_module, random_haiku_module. See random_flax_module doc for an end-to-end example to construct, train, and make prediction with a BNN.
• Wrappers for many TensorFlow Probability distributions in numpyro.contrib.tfp.distributions.
• Wrappers for many TensorFlow MCMC kernels in numpyro.contrib.tfp.mcmc. A user-defined TensorFlow MCMC kernel can be converted to NumPyro-compatible one using TFPKernel.
• New distribution: Geometric distribution.
• New primitive: subsample primitive for data/param subsampling.
• New auto guide: AutoNormal which is similar to AutoDiagonalNormal but more suitable for mean field ELBO and param subsampling.
• New SVI objective: TraceMeanField_ELBO.
• New optimizer: Minimize with BFGS method.

New Examples

• Bayesian Imputation for Missing Values in Discrete Covariates tutorial: leverage enumeration mechanism to marginalize discrete missing covariates - applied for clinical synthesis data.
• Bayesian Hierarchical Linear Regression tutorial: practical Bayesian inference for Kaggle competitions.
• Ordinal Regression tutorial: how to deal with ordered discrete data.

Deprecation

Changes to match Pyro api.

• ELBO objective is renamed to Trace_ELBO.
• value argument in Delta distribution is replaced by v.
• init_strategy argument in autoguides is replaced by init_loc_fn.

Enhancements and Bug Fixes

• Relax simplex constraint. #725 #737
• Fix init_strategy argument not respected in HMC and SA kernels. #728
• Validate the model when cannot find valid initial params. #733
• Avoid nan acceptance probability in SA kernel. #740

Thanks @xidulu, @vanAmsterdam, @TuanNguyen27, @ucals, @elchorro, @RaulPL, and many others for your contributions and helpful feedback!

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!

0.2.4

New Features

• NumPyro can be used on Cloud TPUs.
• Deterministic primitive to record deterministic values in a model.
• Mask handler to mask out the log probability of a sample site using a mask array.
• Sample Adaptive MCMC, a non-gradient based sampler that has a high effective sample size per second.
• New normalizing flow: Block neural autoregressive transform.
• Additional auto guides: AutoLowRankMultivariateNormal and AutoBNAFNormal.

New Examples

• Predator-prey model example: uses MCMC to solve the inverse problem of an ODE system.
• Neural transport example: uses a normalizing flow to transform the posterior to a Gaussian-like one, thereby improving mixing rate for HMC/NUTS.

Deprecation / Breaking Changes

• Predictive's get_samples method is deprecated in favor of __call__ method.
• MCMC constrain_fn is renamed to postprocess_fn.

Enhancements and Bug Fixes

• Change the init scale of Auto*Normal guides from 1. to 0.1 - this is helpful for stability during the early training phase.
• Resolve overflow issue with the Poisson sampler.

0.2.3

Patches 0.2.2 with the following changes:

• restore compatibility with python 3.7 for mcmc.
• impose cache size limit in MCMC utilities.

0.2.2

Breaking changes

• Minor interface changes to MCMC utility functions. All experimental interfaces are marked as such in the documentation.

New Features

• A numpyro.factor primitive that adds an arbitrary log probability factor to a probabilistic model.

Enhancements and Bug Fixes

• Addressed a bug where multiple invocations of MCMC.run would wrongly use the previously cached arguments.
• MCMC reuses compiled model code whenever possible. e.g. when re-running with different but same sized model arguments.
• Ability to reuse the same warmup state for subsequent MCMC runs using MCMC.warmup.