Version 0.6.0

[andsteing]

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

New Contributors

• @sanchit-gandhi made their first contribution in Forward unroll argument in scan_with_axes #2339
• @Chuxiaof made their first contribution in Fix type annotation for step in training.checkpoints #2343
• @Davidnet made their first contribution in Make link a link on the getting started #2340
• @Jeevesh8 made their first contribution in Correct state variable name #2369
• @banda-larga made their first contribution in fixed math expressions  #2392

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

---

This discussion was created from the release Version 0.6.0.