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!