Numpyro

A minor release to fix numpyro 0.13.1 broken on jax 0.4.14

Enhancements and Bug Fixes

- Add promote_batch_shape rule for Independent (1630) by deoxyribose
- Support custom prng key (1642) by fehiepsi

Breaking changes

Drops support for python 3.8 and requires jax version >= 0.4.14

New Features

- Distributions are now vmap-able (1529) (a great contribution by pierreglaser)

Enhancements and Bug Fixes

- Enhance LocScaleReparam's documentation (1599) by Madhav-Kanda
- Fix incorrect unflattenning of inverse transforms (1600) by pierreglaser
- Update Stein mixture (1601 and 1612) by OlaRonning
- Support model without global variables in AutoSemiDAIS (1610 and 1619) by fehiepsi
- Fix mixture assert message: string shoulf be f-string (1617) by adrn
- Add support for local variables in RenyiELBO (1608) by fehiepsi
- Fix quantile computation of mvn autoguides (1622) by fehiepsi
- Respect log_density in kl of delta (1625) by fehiepsi
- Add vectorized_particles to ELBO (1624) by fehiepsi
- Fix bug in SineBivariateVonMises sampler (1628) by deoxyribose

This release is composed of great contributions and feedback from the Pyro community. Thank you!

This release includes a fix for jax 0.4.11 (1595).

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)