Numpyro

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.

Refer to the README for details.