Version 0.6.0

What's Changed

• Add on_commit_callback to put the responsibility of renaming the directories on the users of the serialization library. This will also fix the GCS atomic rename issue where the users can write a success file when the commit is successful and check the existence of that file before deserialization. by @copybara-service in #2328
• RDT Redesign by @cgarciae in #2177
• Fix stale URLs to read the docs site. by @levskaya in #2338
• Replace all jax.tree_* calls with jax.tree_util.tree_* by @levskaya in #2325
• Further fix the singular-leaf checkpointing, and add tests. by @copybara-service in #2336
• Forward unroll argument in scan_with_axes by @sanchit-gandhi in #2339
• Document repo analytics by @cgarciae in #2317
• Improve flax basics by @cgarciae in #2291
• Split build into multiple jobs by @cgarciae in #2277
• Fix type annotation for step in training.checkpoints by @Chuxiaof in #2343
• Add test for writing and restoring empty checkpoints. by @copybara-service in #2345
• Allow all processes to checkpoint when not using GDA by @copybara-service in #2350
• Make importing tensorstore optional and move related type hints to comments. by @copybara-service in #2348
• Use jax.named_scope for name stack rather than named_call. by @copybara-service in #2349
• Internal change by @copybara-service in #2356
• Fix sphinx CI errors by @cgarciae in #2361
• Forward path to rewound Scope by @jheek in #2360
• Make link a link on the getting started by @Davidnet in #2340
• Fix colab & github links by @cgarciae in #2363
• Fix ConvTranspose with circular padding by @cgarciae in #2364
• Add some docstrings to the old flax.training.common_utils module. by @levskaya in #2373
• Correct state variable name by @Jeevesh8 in #2369
• Allow linen's Conv layer to operate on arbitrary-rank inputs. by @copybara-service in #2308
• Copies dynamic_scale.py from optim/ to training/. by @copybara-service in #2375
• Add option auto_flush to flax.metrics.tensorboard.SummaryWriter by @copybara-service in #2376
• updated supported transforms in lifting docs. by @levskaya in #2374
• Test docstrings with autodoc on CI by @cgarciae in #2372
• skip remat test that fails with autodiff by @mattjj in #2389
• Removes flax.optim.dynamic_scale. by @copybara-service in #2314
• Plumb spmd_axis_name from vmap_with_axes through to JAX vmap by @copybara-service in #2390
• fixed math expressions by @banda-larga in #2392

New Contributors

• @sanchit-gandhi made their first contribution in #2339
• @Chuxiaof made their first contribution in #2343
• @Davidnet made their first contribution in #2340
• @Jeevesh8 made their first contribution in #2369
• @banda-larga made their first contribution in #2392

Full Changelog: v0.5.3...v0.6.0