Flax Files

What's Changed

• Add vmap reset() regression test for nnx.metrics (#5483) by @chenkuanliao in https://github.com/google/flax/pull/5491
• Validate dropout_rate in dot_product_attention by @mohsinm-dev in https://github.com/google/flax/pull/5503
• Add RoPE embeddings by @samanklesaria in https://github.com/google/flax/pull/5481
• implement axes_scan's broadcast semantics via dce by @mattjj in https://github.com/google/flax/pull/5540
• Refactor NNX Variable to remove legacy JAX QDD dependencies. by @copybara-service[bot] in https://github.com/google/flax/pull/5541
• [typing] add assertion to avoid pyrefly error by @copybara-service[bot] in https://github.com/google/flax/pull/5546
• add pyrefly suppressions for jax.jit annotation change by @copybara-service[bot] in https://github.com/google/flax/pull/5547
• Early exit in Pytree._setattr for data descriptors. by @copybara-service[bot] in https://github.com/google/flax/pull/5549
• replace lazy_init implementation to avoid jax partial_eval by @mattjj in https://github.com/google/flax/pull/5551
• Make Flax thread-safe for multi-threaded tracing. by @copybara-service[bot] in https://github.com/google/flax/pull/5543
• release Flax 0.12.9 by @copybara-service[bot] in https://github.com/google/flax/pull/5557

New Contributors

• @chenkuanliao made their first contribution in https://github.com/google/flax/pull/5491

Full Changelog: https://github.com/google/flax/compare/v0.12.8...v0.12.9