0.12.0

New Features

• New distribution: Gompertz distribution (#1551)
• New initialization strategy: 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 (#1514)
  • A port of Gaussian Mixture Model tutorial from Pyro (#1562)
  • A toy mixture model with discrete enumeration (#1568)
• New inference utilities get_transforms and 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)