0.21.0

What's Changed

• Add pupport forward-mode differentiation with autoguide by @juanitorduz in #2163
• Fix TraceEnum_ELBO ignoring is_auxiliary option by @juanitorduz in #2164
• Fix: Tracer leaks during NUTS sampling on random_nnx_module when NNX module contains mutable state by @juanitorduz in #2162
• FIX: Running MCMC within JIT causes tracer leaks by @juanitorduz in #2165
• Example: Embarrassingly Parallel MCMC via Consensus Monte Carlo by @juanitorduz in #2167
• Remove vectorize_particles from RenyiELBO docstring (Issue #2166) by @juanitorduz in #2168
• fix repeated recompilation in no-progress-bar fori_collect path by @simeoncampos in #2171
• Update GPJax notebook to simplified API by @thomaspinder in #2177
• Effect handlers Tutorial by @juanitorduz in #2169
• Generate outputs by @thomaspinder in #2179
• improve type error message for 'fn' in sample function by @Qazalbash in #2175
• Fix xla_pmap_p import for JAX versions that removed pmap by @saitcakmak in #2173
• Change key type from prng_key to Array by @fehiepsi in #2178
• feat(gh-2142): add VSBC diagnostic by @michaelellis003 in #2170
• Allow zero rate for Poisson distribution. by @tillahoffmann in #2176
• Bump to 0.21.0 by @fehiepsi in #2180

New Contributors

• @simeoncampos made their first contribution in #2171
• @saitcakmak made their first contribution in #2173

Full Changelog: 0.20.1...0.21.0