Numpyro

Breaking Changes

- HMC's [find_heuristic_step_size](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.hmc.HMC) (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](https://github.com/pyro-ppl/numpyro/#a-simple-example---8-schools) for the new usage pattern.
- [Automatic Guide Generation](http://num.pyro.ai/en/latest/autoguide.html) module is moved from `numpyro.contrib.autoguide` to the main inference module `numpyro.infer.autoguide`.
- Various API changes to match Pyro API:
+ [mask](http://num.pyro.ai/en/latest/handlers.html#mask) handler: `mask_array` arg is renamed to `mask`.
+ [scale](http://num.pyro.ai/en/latest/handlers.html#scale) handler: `scale_factor` arg is renamed to `scale`.
+ [condition](http://num.pyro.ai/en/latest/handlers.html#condition) and [substitute](http://num.pyro.ai/en/latest/handlers.html#substitute) handlers: `param_map` is renamed to `data`.
+ `MultivariateAffineTransform` transform is renamed to [LowerCholeskyAffine](http://num.pyro.ai/en/latest/distributions.html#lowercholeskyaffine).
+ `init_to_prior` strategy is renamed to [init_to_sample](http://num.pyro.ai/en/latest/utilities.html#init-to-sample).

New Features

- [Funsor-based NumPyro](http://num.pyro.ai/en/latest/funsor.html): allow enumeration over discrete latent variables. See mixture and Markov examples below for some applications.
- New primitives: [plate_stack](http://num.pyro.ai/en/latest/primitives.html#numpyro.primitives.plate_stack) and [scan](http://num.pyro.ai/en/latest/primitives.html#numpyro.contrib.control_flow.scan). If your model has Python `for` loop, considering using `scan` instead to improve compiling time.
- New handlers: [reparam](http://num.pyro.ai/en/latest/handlers.html#reparam), [scope](http://num.pyro.ai/en/latest/handlers.html#scope), and [lift](http://num.pyro.ai/en/latest/handlers.html#lift).
- New distributions: [von Mises](http://num.pyro.ai/en/latest/distributions.html#vonmises), [Gumbel](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.continuous.Gumbel), [Logistic](http://num.pyro.ai/en/latest/distributions.html#logistic), [Laplace](http://num.pyro.ai/en/latest/distributions.html#laplace), [TruncatedPolyaGamma](http://num.pyro.ai/en/latest/distributions.html#truncatedpolyagamma), [ExpandedDistribution](http://num.pyro.ai/en/latest/distributions.html#expandeddistribution), [MaskedDistribution](http://num.pyro.ai/en/latest/distributions.html#maskeddistribution), and [ImproperUniform](http://num.pyro.ai/en/latest/distributions.html#improperuniform).
- [Distribution](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution) has new properties `is_discrete`, `has_enumerate_support`, and new methods [shape](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.shape), [enumerate_support](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.enumerate_support), [expand](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.expand), [expand_by](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.expand_by), [mask](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.mask). In addition, `Distribution` has been registered as [a JAX Pytree](https://jax.readthedocs.io/en/latest/pytrees.html) class, with corresponding methods `tree_flatten` and `tree_unflatten`.
- New constraint: [less_than](http://num.pyro.ai/en/latest/distributions.html#less-than).
- Port [Tensor Indexing](http://num.pyro.ai/en/latest/utilities.html#module-numpyro.contrib.indexing) from Pyro.
- Port some [Reparameterizers](http://num.pyro.ai/en/latest/reparam.html) from Pyro.
- Add `batch_ndims` arg to [Predictive](http://num.pyro.ai/en/latest/utilities.html#numpyro.infer.util.Predictive) and [log_likelihood](http://num.pyro.ai/en/latest/utilities.html#numpyro.infer.util.log_likelihood) to allow using those utilities with arbitrary number of batch dimensions.

New Examples

- [Proportion Test example](http://pyro.ai/numpyro/examples/proportion_test.html): 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](http://pyro.ai/numpyro/examples/annotation.html): illustrates enumeration for mixture models.
- [Enumerate HMM examples](http://pyro.ai/numpyro/examples/hmm_enum.html): illustrates enumeration for markov models.
- [Bayesian Imputation tutorial](http://pyro.ai/numpyro/bayesian_imputation.html).

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](http://pyro.ai/numpyro/time_series_forecasting.html) 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](http://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.mcmc.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!

New Features

- NumPyro can be used on Cloud TPUs.
- [Deterministic](http://num.pyro.ai/en/latest/primitives.html#deterministic) primitive to record deterministic values in a model.
- [Mask](http://num.pyro.ai/en/latest/handlers.html#mask) handler to mask out the log probability of a sample site using a mask array.
- [Sample Adaptive MCMC](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.SA), a non-gradient based sampler that has a high effective sample size per second.
- New normalizing flow: [Block neural autoregressive transform](http://num.pyro.ai/en/latest/distributions.html#blockneuralautoregressivetransform).
- Additional auto guides: [AutoLowRankMultivariateNormal](http://num.pyro.ai/en/latest/autoguide.html#autolowrankmultivariatenormal) and [AutoBNAFNormal](http://num.pyro.ai/en/latest/autoguide.html#autobnafnormal).

New Examples

- [Predator-prey model example](https://github.com/pyro-ppl/numpyro/blob/master/examples/ode.py): uses MCMC to solve the inverse problem of an ODE system.
- [Neural transport example](https://github.com/pyro-ppl/numpyro/blob/master/examples/neutra.py): 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](http://num.pyro.ai/en/latest/utilities.html#numpyro.infer.util.Predictive.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.

Patches `0.2.2` with the following changes:
- restore compatibility with python 3.7 for `mcmc`.
- impose cache size limit in MCMC utilities.

Breaking changes

- Minor interface changes to [MCMC utility](http://num.pyro.ai/en/latest/utilities.html#module-numpyro.infer.util) functions. All experimental interfaces are marked as such in the documentation.

New Features

- A [numpyro.factor](http://num.pyro.ai/en/latest/primitives.html#numpyro.primitives.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](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.MCMC.warmup).

Breaking changes

- Code reorganization - `numpyro.mcmc` is moved to `numpyro.infer.mcmc` but all major classes are exposed in the `numpyro.infer` module.
- `rng` argument to many classes and the `seed` handler has been more accurately renamed to `rng_key`.
- Deprecated functions that formed the old interface like `mcmc` and `svi` have been removed.

New Features

- Improved turning condition for NUTS that results in much higher effective sample size for many models.
- A [numpyro.plate](http://num.pyro.ai/en/latest/primitives.html#plate) context manager, which records conditional independence information in the trace and does automatic broadcasting, like in Pyro.
- Inclusion of [AutoMultivariateNormal](http://num.pyro.ai/en/latest/autoguide.html#automultivariatenormal), [AutoLaplaceApproximation](http://num.pyro.ai/en/latest/autoguide.html#autolaplaceapproximation) to the [autoguide](http://num.pyro.ai/en/latest/autoguide.html) module.
- More distributions like [LowRankMultivariateNormal](http://num.pyro.ai/en/latest/distributions.html#lowrankmultivariatenormal), [LKJ](http://num.pyro.ai/en/latest/distributions.html#lkj), [BetaBinomial](http://num.pyro.ai/en/latest/distributions.html#betabinomial), [GammaPoisson](http://num.pyro.ai/en/latest/distributions.html#gammapoisson), [ZeroInflatedPoisson](http://num.pyro.ai/en/latest/distributions.html#zeroinflatedpoisson), and [OrderedLogistic](http://num.pyro.ai/en/latest/distributions.html#orderedlogistic).
- More transforms: [MultivariateAffineTransform](http://num.pyro.ai/en/latest/distributions.html#multivariateaffinetransform), [InvCholeskyTransform](http://num.pyro.ai/en/latest/distributions.html#invcholeskytransform), [OrderedTransform](http://num.pyro.ai/en/latest/distributions.html#orderedtransform).
- A `numpyro.compat` module that supports the [pyro generic](https://github.com/pyro-ppl/pyro-api) API for modeling and inference that can dispatch to multiple Pyro backends.
- Inclusion of [Independent](http://num.pyro.ai/en/latest/distributions.html#independent) distribution and [`Distribution.to_event`](http://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.distribution.Distribution.to_event) method to convert independent batch dimensions to dependent event dimensions. See the Pyro tutorial on [tensor shapes](http://pyro.ai/examples/tensor_shapes.html#Reshaping-distributions) for more details.
- A [Predictive](http://num.pyro.ai/en/latest/utilities.html#predictive) utility for generating samples from prior models, predictions from models using SVI's guide, or posterior samples from MCMC.
- A [log_likelihood](http://num.pyro.ai/en/latest/utilities.html#log-likelihood) utility function that can compute the log likelihood for observed data by conditioning latent sites to values from the posterior distribution.
- New [ClippedAdam](http://num.pyro.ai/en/latest/optimizers.html#clippedadam) optimizer to prevent exploding gradients.
- New [RenyiELBO](http://num.pyro.ai/en/latest/svi.html#renyielbo) loss for Renyi divergence variational inference and importance weighted variational inference.

Enhancements and Bug Fixes

- MCMC does not throw an error on models with no latent sites.
- [numpyro.seed](http://num.pyro.ai/en/latest/handlers.html#numpyro.handlers.seed) handler can be used as a context manager like:
python
with numpyro.seed(rng_seed=1):
...

- Utilities to enable [validation checks](http://num.pyro.ai/en/latest/utilities.html#numpyro.distributions.distribution.enable_validation) for distributions, set [host device count](http://num.pyro.ai/en/latest/utilities.html#set-host-device-count), and [platform](http://num.pyro.ai/en/latest/utilities.html#set-platform).
- More efficient sampling from Binomial / Multinomial distributions.
- The evidence lower bound loss for SVI is now a class called `ELBO`.
- Add `energy` field to [HMCState](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.HMCState), which is used to compute [Bayesian Fraction of Missing Information](https://mc-stan.org/misc/warnings.html#bfmi-low) for diagnostics.
- Add `init_strategy` arg to [HMC](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.HMC)/[NUTS](http://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.NUTS) classes, which allows users select various [initialization strategies](http://num.pyro.ai/en/latest/utilities.html#initialization-strategies).

Highlights

- **Interface Changes to MCMC and SVI**: The interface for inference algorithms have been simplified, and is much closer to Pyro. See [MCMC](https://numpyro.readthedocs.io/en/stable/mcmc.html#numpyro.mcmc.MCMC) and [SVI](https://numpyro.readthedocs.io/en/stable/svi.html#numpyro.svi.SVI).
- **Multi-chain Sampling for MCMC**: There are three options provided: `parallel` (default), `sequential`, and `vectorized`. Currently, `parallel` method is the fastest among the three.

Breaking changes

- The primitives `param`, `sample` are moved to [primitives](https://numpyro.readthedocs.io/en/stable/primitives.html) module. All primities are exposed in `numpyro` namespace.

New Features

MCMC
- In MCMC, we have the option to collect fields other than just the samples such as number of steps or step size, using `collect_fields` arg in [MCMC.run](https://numpyro.readthedocs.io/en/stable/mcmc.html#numpyro.mcmc.MCMC.run). This can be useful when gathering diagnostic information during debugging.
- `diverging` field is added to [HMCState](https://numpyro.readthedocs.io/en/stable/mcmc.html#numpyro.mcmc.HMCState). This field is useful to detect divergent transitions.
- Support improper prior through `param` primitives. e.g.
python
def model(data):
loc = numpyro.param('loc', 0.)
scale = numpyro.param('scale', 0.5, constraint=constraints.positive)
return numpyro.sample('obs', dist.Normal(loc, scale), obs=data)


Primitives / Effect Handlers
- [module](https://numpyro.readthedocs.io/en/stable/primitives.html#module) primitive to support JAX style neural network. See [VAE example](http://pyro.ai/numpyro/vae.html).
- [condition](https://numpyro.readthedocs.io/en/stable/handlers.html#condition) handler for conditioning sample sites to observed data.
- [scale](https://numpyro.readthedocs.io/en/stable/handlers.html#scale) handler for rescaling the log probability score.

Optimizers

JAX optimizers are wrapped in the [numpyro.optim](https://numpyro.readthedocs.io/en/stable/optimizers.html#module-numpyro.optim) module, so that the optimizers can be passed in directly to `SVI`.

Distributions
- New distributions: Delta, GaussianRandomWalk, InverseGamma, [LKJCholesky](https://numpyro.readthedocs.io/en/stable/distributions.html#lkjcholesky) (with both `cvine` and `onion` methods for sampling), MultivariateNormal.
- New transforms: [CorrCholeskyTransform](https://numpyro.readthedocs.io/en/stable/distributions.html#corrcholeskytransform) (which is vectorized), [InverseAutoregressiveTransform](https://numpyro.readthedocs.io/en/stable/distributions.html#numpyro.distributions.flows.InverseAutoregressiveTransform), LowerCholeskyTransform, PermuteTransform, PowerTransform.

Utilities

- [predictive](https://numpyro.readthedocs.io/en/stable/utilities.html#predictive) utility for vectorized predictions from the posterior predictive distribution.

Autoguides

An experimental [autoguide](https://numpyro.readthedocs.io/en/stable/autoguide.html) module, with more autoguides to come.

New Examples

- [Sparse Linear Regression](http://pyro.ai/numpyro/sparse_regression.html) - fast Bayesian discovery of pairwise interactions in high dimensional data.
- [Gaussian Process](http://pyro.ai/numpyro/gp.html) - sample from the posterior over the hyperparameters of a gaussian process.
- [HMC on Neal's Funnel](http://pyro.ai/numpyro/funnel.html) - automatic reparameterization through transform distributions.

Enhancements and Bug Fixes
- Improve compiling time in MCMC.
- Better PRNG splitting mechanism in SVI (to avoid reusing PRNG keys).
- Correctly handle models with dynamically changing distribution constraints. e.g.
python
def model():
x = numpyro.sample('x', dist.Uniform(0., 2.))
y = numpyro.sample('y', dist.Uniform(0., x)) y's support is not static.

- Fixes `step_size` getting `NaN` in MCMC when it becomes extremely small.