Version 0.6.1

What's Changed

• Updates examples/{imagenet,wmt} requirements. by @andsteing in #2405
• Bump rich dependency version by @yklcs in #2407
• Adds axis_name and axis_index_groups to LayerNorm and GroupNorm. by @copybara-service in #2402
• Plumb spmd_axis_name through transforms.vmap through to JAX vmap by @copybara-service in #2398
• Support multiple inputs in flax lifted vjp/custom_vjp by @copybara-service in #2399
• Explicit reexport initializers from jax by @lkhphuc in #2409
• Improve tabulate by @cgarciae in #2316
• Add path_aware_map function by @cgarciae in #2371
• PIL does not accept DeviceArray, so needed to use numpy. by @villeh1 in #2427
• Move examples to RTD by @cgarciae in #2367
• Simplify dynamic context by @cgarciae in #2388
• Remove pytype generic workaround by @jheek in #2446
• Add static_argnums to nn.checkpoint by @cgarciae in #2457
• ignore tf deprecation warning. by @copybara-service in #2466
• Fix Managing Parameters and State docs by @cgarciae in #2473
• Use gfile.listdir instead of gfile.glob in checkpointing by @IvyZX in #2470
• Create test matrix to speedup tests by @cgarciae in #2458
• Fix Conv docstrings by @cgarciae in #2425
• Use proper scikit-learn dependency by @cgarciae in #2465
• Improve attribute error msg for unbounded modules by @cgarciae in #2440
• Adding "count_include_pad" argument to flax.linen.pooling.avg_pool by @dslisleedh in #2451
• Add perturb() to allow capturing intermediate gradients by @IvyZX in #2476
• fix DynamicScale docstring by @cgarciae in #2491
• test against python 3.8 and 3.9 by @cgarciae in #2490
• Update version to 0.6.1 by @cgarciae in #2494
• Adoption cache should use WeakValueDictionary. by @levskaya in #2495
• FLIP: General metadata by @jheek in #2435

New Contributors

• @yklcs made their first contribution in #2407
• @villeh1 made their first contribution in #2427
• @dslisleedh made their first contribution in #2451

Full Changelog: v0.6.0...v0.6.1