Optax 0.2.9

What's Changed

• Start of 0.2.9 development. by @copybara-service[bot] in #1638
• Minor improvements and cleanup using tree utility functions. by @carlosgmartin in #1365
• [JAX] Suppress type errors found by pytype after correcting definition of jax.typing.ArrayLike. by @copybara-service[bot] in #1648
• Update microbatching to support size 0 batches. This change is important for Jax privacy, where the batch size is random, and possibly empty. by @copybara-service[bot] in #1647
• Fix version checks in microbatching by @copybara-service[bot] in #1655
• sigmoid_binary_cross_entropy: add type annotations by @anshuljoshi in #1653
• Remove TensorFlow dependency in the differentially private CNN on MNIST example by @rajasekharporeddy in #1610
• Temporary patch for failing test on lars on complex numbers with gpu by @copybara-service[bot] in #1661
• Avoid unnecessary directory change in test.sh. by @carlosgmartin in #1324
• Muon: fix excessive memory from Newton-Schulz unroll=True by @bzantium in #1650
• Automated Code Change by @copybara-service[bot] in #1663
• Migrate from flax.linen to flax.nnx in Lookahead Optimizer in MNIST example. by @rajasekharporeddy in #1589
• Remove TensorFlow dependency in Reduce on Plateau Learning Rate Scheduler example by @rajasekharporeddy in #1616
• Update mlp_mnist notebook for removing tf dependency and NNX migration by @selamw1 in #1625
• Remove TensorFlow dependency and migrate cifar10_resnet to Flax NNX by @selamw1 in #1508
• Improve microbatching support for positional args passed by name and other variations of positional/kwargs. by @copybara-service[bot] in #1665
• Add ArrayTree to the public API of optax. by @copybara-service[bot] in #1666
• Improve microbatching support for positional args passed by name and other variations of positional/kwargs. by @copybara-service[bot] in #1668
• Improve robustness of microbatching w.r.t. different calling conventions (Args / kwargs). by @copybara-service[bot] in #1669
• Fix CI by relaxing tolerance in optax nnls test by @copybara-service[bot] in #1670
• Remove the stochastic gradient estimatorsand control variates documentation by @rajasekharporeddy in #1664
• Remove non-monotonicity check from test_nnls as it is theoretically not guaranteed. by @copybara-service[bot] in #1672
• Note hyperparams can be strongly typed JAX arrays. by @copybara-service[bot] in #1673
• Remove packaging dependency and vendor version parsing in Optax. by @copybara-service[bot] in #1675
• Split L-BFGS tests out of alias_test.py into lbfgs_test.py. by @copybara-service[bot] in #1676
• Migrate from pytype to pyrefly for type checking. by @copybara-service[bot] in #1674
• Fix pyrefly type-checking failures in GitHub CI and simplify freezing tests. by @copybara-service[bot] in #1682
• Fix wrong schedule name in cosine_onecycle_schedule error message by @adityasingh2400 in #1680
• feat: add optional weight_decay to adabelief optimizer by @yasumorishima in #1634
• Repair use of deprecated parameter by @NeilGirdhar in #1671
• Removing pytype or pyrefly from tests. by @copybara-service[bot] in #1684
• Improve microbatch and micro_grad documentation. by @copybara-service[bot] in #1685
• docs: render one page per public symbol via autosummary toctree by @adityasingh2400 in #1679
• relax tolerance on the vdot test by @copybara-service[bot] in #1691
• fix: prevent PyPI/conda-forge dependency drift with automated release checks by @Muneerali199 in #1688
• fix documentation of microbatch. micro_grad actually computes per-example gradients, i.e., microbatch(vmap(grad(fun))), which is a bit different than microbatch(grad(fun)). by @copybara-service[bot] in #1694
• Reduce focal loss memory usage. by @copybara-service[bot] in #1700
• Migrate to the latest flax.nnx module from the old flax.linen module in Gradient Accumulation example. by @rajasekharporeddy in #1699
• Fix mlc_config.json for markdown-link-check by @copybara-service[bot] in #1705
• docs(losses): clarify cosine_similarity/cosine_distance epsilon clips the squared norm by @Osamaali313 in #1703
• Migrate from the old flax.linen to the latest flax.nnx module in Simple NN with Flax example. by @rajasekharporeddy in #1698
• Ignore badge in link checks by @copybara-service[bot] in #1707
• Adding mask for weight decay in Prodigy-AdamW implementation by @copybara-service[bot] in #1702
• Resolve various type checking and indexing warnings in core optimizers and losses. by @copybara-service[bot] in #1677
• Resolve various type checking and indexing warnings in core optimizers and losses. by @copybara-service[bot] in #1719
• contrib: add sharding tests for contrib optimizers by @irhyl in #1714
• Set maximum_signature_line_length = 1 in docs/conf.py. by @carlosgmartin in #1723
• Avoid stateful directory changes by using make command's -C flag. by @carlosgmartin in #1722
• Fixes a TypeError in rosenbrock_ademamix.ipynb caused by a parameter name mismatch. by @rajasekharporeddy in #1741
• add pyrefly suppressions for jax.jit annotation change by @copybara-service[bot] in #1748
• Fix the Attribute Error in Differentially private CNN on MNIST example notebook by @rajasekharporeddy in #1739
• Fix AttributeError in Sharpness-Aware Minimization (SAM) example notebook by using _replace for NamedTuple by @rajasekharporeddy in #1742
• Migrate the Adversarial training example to the latest flax.nnx module. by @rajasekharporeddy in #1710
• docs(microbatching): Fix docstring formatting and add Sphinx cross-references by @rajasekharporeddy in #1758
• Fix the Attribute Error in gradient_accumulation_and_microbatching.ipynb example notebook by @rajasekharporeddy in #1740
• Fix pyrefly type-check failures in the OSS CI. by @copybara-service[bot] in #1764
• Add weight_decay_mask to dadapt_adamw, and fix it for prodigy. by @copybara-service[bot] in #1763
• Migrate Reduce on Plateau Learning Rate Scheduler example to the latest flax.nnx module by @rajasekharporeddy in #1711
• Fix OSS pyrefly issues in linear_algebra, classification losses, and state_utils. by @copybara-service[bot] in #1779
• Correctly resolve muon precision tests on accelerators by @copybara-service[bot] in #1778
• Fix type errors reported by pyrefly 1.3.0 by @copybara-service[bot] in #1777
• Remove leftover pytype comments across the codebase. by @copybara-service[bot] in #1782
• Fix AdamW example in getting started notebook by @rishabhsai in #1780
• Fix negative CTC loss for confident alignments in float32 by @rajasekharporeddy in #1772
• docs: fix typos and update docstring formatting across multiple modules by @rajasekharporeddy in #1759
• Automated Code Change by @copybara-service[bot] in #1807
• [XLA:GPU] Configure test matmul precision in Optax contrib common test for BF16 default by @copybara-service[bot] in #1802
• Migrate to the latest flax.nnx module from the legacy flax.linen module in Character-level Transformer on Tiny Shakespeare example. by @rajasekharporeddy in #1709
• Mask cosine inputs before normalization to preserve finite gradients by @sylvesterkaczmarek in #1801
• docs(contrib): fix paper reference link for ADOPT optimizer by @himoto in #1731
• fix(clipping): keep zero-update reverse gradients finite by @jabrailkhalil in #1815
• Initialize the first moment of scale_by_yogi to zero by @satyam-mishra-dev in #1818
• Raise a ValueError at trace time when eps underflows in the optimizer state dtype by @copybara-service[bot] in #1835
• • Compute ntxent pairwise cosine similarities in float32/float64 via jax.lax.rsqrt normalization and 2D jnp.matmul (O(N^2 D) on the MXU, using Precision.HIGHEST on float32/float64 embeddings to avoid bfloat16/TF32 operand truncation) instead of 3D [N, N, D] broadcast cosine_similarity + [N, 2N] concatenation. by @copybara-service[bot] in #1848
• Remove unused legacy safe_norm alias from optax._src.utils. by @copybara-service[bot] in #1869
• Release Optax version 0.2.9. by @copybara-service[bot] in #1875

New Contributors

• @anshuljoshi made their first contribution in #1653
• @bzantium made their first contribution in #1650
• @adityasingh2400 made their first contribution in #1680
• @yasumorishima made their first contribution in #1634
• @Muneerali199 made their first contribution in #1688
• @Osamaali313 made their first contribution in #1703
• @irhyl made their first contribution in #1714
• @rishabhsai made their first contribution in #1780
• @sylvesterkaczmarek made their first contribution in #1801
• @himoto made their first contribution in #1731
• @jabrailkhalil made their first contribution in #1815
• @satyam-mishra-dev made their first contribution in #1818

Full Changelog: v0.2.8...v0.2.9