Numpyro

New Features

- New distribution: [Gompertz distribution](https://num.pyro.ai/en/latest/distributions.html#gompertz) (1551)
- New initialization strategy: [init_to_mean](https://num.pyro.ai/en/latest/utilities.html#init-to-mean) (1550)
- New examples and tutorials:
+ Illustrate the usage of JAX PositionalSharding for distributing the computations of log_prob/grad over multiple devices in [MCMC](https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.mcmc.MCMC) (1514)
+ A port of [Gaussian Mixture Model](https://num.pyro.ai/en/latest/tutorials/gmm.html) tutorial from Pyro (#1562)
+ [A toy mixture model with discrete enumeration](https://num.pyro.ai/en/latest/examples/toy_mixture_model_discrete_enumeration.html) (#1568)
- New inference utilities [get_transforms](https://num.pyro.ai/en/latest/utilities.html#get-transforms) and [unconstrain_fn](https://num.pyro.ai/en/latest/utilities.html#unconstrain-fn) to transform between unconstrained and constrained space (1564)
- Support jaxns>=2.0.1 (1546)

Enhancements and Bug Fixes
- Make transforms jittable (1575)
- Fixed typo in surrogate posterior of beta (1591)
- Do not scale mnist label (1589)
- Do not mutate shapes of ExpandedDistribution for map-free ops (1574)
- Add support for JAX custom PRNG (1587)
- Include deterministic variables in AutoDelta's sample_posterior (1584)
- Fix forward shape of SimplexToOrderTransform (1583)
- Fix inf's in TruncatedNormal log_prob & sample (1581)
- Allow users to specify total_count_max in Multinomial (1557)
- Allow pickled mcmc object to run post warmup phase (1558)
- Add init_params argument to svi.init() and svi.run() (1561)
- Support pickling MCMC objects with enumeration (1577)
- Raise error when reparameterize lognormal (1548)
- Avoid initializing model params when already specified in guide (1553)
- Respect init params if provided to mcmc.run (1547)
- Fix provenance for jax 0.4.4 (1543)
- Use analytic kl divergence in TraceEnum_ELBO (1533)
- Properly handle contraction of guide plates in TraceEnum_ELBO (1537)
- Raise an error if there is no common scale when model enumerated (1536)
- Optimize reduction of enumerated guide sites (1531)
- Guess max_plate_nesting in TraceEnum_ELBO (1528)

Breaking changes

Drop Python 3.7 support and require the minimal jax version 0.4.

New Features

- New distributions:
+ [EulerMaruyama](https://num.pyro.ai/en/latest/distributions.html#eulermaruyama) for modelling stochastic differential equations (SDE) (thanks to yayami3)
+ [GaussianCopula](https://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.copula.GaussianCopula) and [GaussianCopulaBeta](https://num.pyro.ai/en/latest/distributions.html#gaussiancopulabeta) (i.e. MultivariateBeta) to incorporate correlations into univariate random variables (thanks to hessammehr)
+ [MatrixNormal](https://num.pyro.ai/en/latest/distributions.html#matrixnormal) generalizes multivariate normal to matrix-valued random variables (thanks to kaijennissen)
+ [LogUniform](https://num.pyro.ai/en/latest/distributions.html#loguniform) distribution (thanks to yayami3 and andrewfowlie)
- New objective [TraceEnum_ELBO](https://num.pyro.ai/en/stable/svi.html#numpyro.infer.elbo.TraceEnum_ELBO) to support enumeration in SVI (thanks to ordabayevy)
- New tutorial:
+ [Text-Based Ideal Points using NumPyro](https://num.pyro.ai/en/latest/tutorials/tbip.html#Text-Based-Ideal-Points-using-NumPyro) (thanks to elchorro)
- 1508 Allow [rendering deterministic sites](https://num.pyro.ai/en/latest/tutorials/model_rendering.html#Rendering-deterministic-sites) (thanks to YanniPapandreou)

Enhancements and Bug Fixes

- 1507 Made constraints robust to pickling (thanks to pierreglaser)
- 1515 Fixed log_prob for negative correlation in SineBivariateVonMises distribution (thanks to OlaRonning)
- 1509 Fixed a bug at flatten/unflatten distributions which mixed the order of arguments of the distributions (thanks to hessammehr)
- 1494 Fixed Mixture distribution with unnormalized logits giving incorrect log_prob
- 1486 Returned a positive ordered vector when applies ExpTransform for an ordered vector
- 1491 Fixed Mixture intermediate values
- 1480 Fixed some computations in Bayesian Hierarchical Stacking tutorial (thanks to cpieringer)
- 1478 Added icdf methods for Beta, Gamma, StudentT
- 1477 Allowed multiple arguments to initialize flax/haiku modules (thanks to fehiepsi)
- 1475 Used TFP's `betainc` (which supports taking gradient w.r.t. parameters) in StudentT.cdf (thanks to colehaus)

New Features

- 1458 New [Mixture](https://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.mixtures.Mixture) distribution to allow component distributions come from different families (thanks to dfm)

Enhancements and Bug Fixes

- 1469 Fix memory leak when using scan
- 1465 Fix validation logic for transformed distribution (thanks to xidulu)
- 1462 Avoid infinite loop in vmapped Binomial (thanks to tbenthompson)
- 1449 Move `contrib.render` functionality to `infer.inspect` (thanks to mattja)

New Features

- 1394 New distribution Conditional Autoregressive [CAR](https://num.pyro.ai/en/latest/distributions.html?#car) (thanks to theorashid)
- 1434 New flexible auto guides for models with both global and local latent variables: [AutoSemiDAIS](https://num.pyro.ai/en/latest/autoguide.html#numpyro.infer.autoguide.AutoSemiDAIS) and [AutoSurrogateLikelihoodDAIS](https://num.pyro.ai/en/latest/autoguide.html#numpyro.infer.autoguide.AutoSurrogateLikelihoodDAIS)
- 1429 New example: [Conditional Variational Autoencoder in Flax](https://num.pyro.ai/en/latest/examples/cvae.html) (thanks to dirmeier)

Enhancements and Bug Fixes

- 1401 Fix `obs` argument is not respected when `sample` primitive is not executed under any handler (thanks to gcskoenig)
- 1412 TraceGraph_ELBO implementation using provenance tracking
- 1418 Fix SA sampler cannot be run in parallel chains
- 1419 Fix categorical sampler occasionally generate out-of-support samples
- 1436 Allow to use `potential_fn` in BarkerMH
- 1437 Fix for AutoMultivariateNormal.get_posterior method return incorrect distribution (thanks to xidulu)
- 1444 Promote shapes for observed variables inside `scan`'s transition function
- 1443 Consider the time dimension of markov models with `history=0` as plate
- 1441 More stable check for corr_cholesky constraint
- 1400 Fix SineBivariateVonMises sampler

New Features

- 1381 New `render_params` argument for [render_model](https://num.pyro.ai/en/latest/utilities.html#render-model)
- 1366 Allow a fixed number of steps in HMC.
- New tutorials and examples:
+ [Modelling mortality over space and time](https://num.pyro.ai/en/latest/examples/mortality.html)
+ New section on adding intermediate levels to [Bayesian Hierarchical Linear Regression](https://num.pyro.ai/en/latest/tutorials/bayesian_hierarchical_linear_regression.html#5.-Add-layer-to-model-hierarchy:-Smoking-Status) tutorial

Enhancements and Bug Fixes

- 1386 Support JAX 0.3.5
- 1388 Update jaxns to 1.0.0
- 1372 Fix batch shapes of `SineBivariateVonMises` distribution
- 1375 Add `__repr__` method for constraints
- 1358 Force validate_args to be keyword argument
- 1350 Enhance the AR2 example

Thanks, cstoafer, hesenp, tcbegley, themrzmaster, karm-patel, theorashid

New Features

- New distributions: [AsymmetricLaplace](https://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.continuous.AsymmetricLaplace) and [AsymmetricLaplaceQuantile](https://num.pyro.ai/en/latest/distributions.html#numpyro.distributions.continuous.AsymmetricLaplaceQuantile) for quantile regression.
- New tutorials and examples:
+ [Bayesian neural network with SteinVI](https://num.pyro.ai/en/latest/examples/stein_bnn.html)
+ [Deep markov model inferred using SteinVI](https://num.pyro.ai/en/latest/examples/stein_dmm.html)
+ [Zero-Inflated Poisson regression model](https://num.pyro.ai/en/latest/examples/zero_inflated_poisson.html)

Enhancements and Bug Fixes

- 1320 Fix provenance logic for `numpyro.render_model`
- 1330 Fix vectorize sampling for sites with size=0
- 1331 Allow nested params in SteinVI
- 1333 Reduce memory consumption for SteinVI
- 1325 Fix compat MCMC signature

Thanks, Vinnie-Palazeti, wataruhashimoto52, hessammehr, OlaRonning, d-diaz!