Flax Files

What's Changed

• Annotate wrapped_fn's self argument in decorator_lift_transform for trace-ability by @copybara-service[bot] in https://github.com/google/flax/pull/5298
• add with_attributes by @copybara-service[bot] in https://github.com/google/flax/pull/5308
• Remove optax pin by @samanklesaria in https://github.com/google/flax/pull/5292
• Removed redundant code checking wrt in nnx.Optimizer by @vfdev-5 in https://github.com/google/flax/pull/5226
• Fix PyTreeNode + Generic losing parameters when Generic is last in bases by @mohsinm-dev in https://github.com/google/flax/pull/5237
• add recursive_map example by @copybara-service[bot] in https://github.com/google/flax/pull/5311
• Add sharding propagation support in nnx.eval_shape (clone of [#5111]) by @samanklesaria in https://github.com/google/flax/pull/5247
• fix eval_shape's _to_variable by @copybara-service[bot] in https://github.com/google/flax/pull/5316
• support out_sharding mapping in set_metadata by @copybara-service[bot] in https://github.com/google/flax/pull/5312
• feat(nnx): add GQA support to MultiHeadAttention by @ayulockedin in https://github.com/google/flax/pull/5259
• Allow DenyList to be compared. by @copybara-service[bot] in https://github.com/google/flax/pull/5322
• add graph_updates argument to jit by @copybara-service[bot] in https://github.com/google/flax/pull/5317
• add graph_updates argument in shard_map by @copybara-service[bot] in https://github.com/google/flax/pull/5319
• Remove jax/tools/colab_tpu.py. by @copybara-service[bot] in https://github.com/google/flax/pull/5324
• add graph_updates to vmap by @copybara-service[bot] in https://github.com/google/flax/pull/5320
• Add split method to RngStream by @samanklesaria in https://github.com/google/flax/pull/5270
• add graph_updates for scan by @copybara-service[bot] in https://github.com/google/flax/pull/5327
• add graph_updates to while_loop by @copybara-service[bot] in https://github.com/google/flax/pull/5328
• Added support for data masking in Average, Accuracy and MultiMetric by @vfdev-5 in https://github.com/google/flax/pull/5326
• add graph_updates to fori_loop by @copybara-service[bot] in https://github.com/google/flax/pull/5329
• add graph_updates to pmap, grad, and value_and_grad by @copybara-service[bot] in https://github.com/google/flax/pull/5330
• add tree-mode-nnx FLIP by @copybara-service[bot] in https://github.com/google/flax/pull/5310
• Copybara import of the project: by @copybara-service[bot] in https://github.com/google/flax/pull/5331
• add graph_updates to remat by @copybara-service[bot] in https://github.com/google/flax/pull/5336
• add graph_updates to eval_shape and checkify by @copybara-service[bot] in https://github.com/google/flax/pull/5338
• add compat module by @copybara-service[bot] in https://github.com/google/flax/pull/5340
• add more tests that check for consistent aliasing in transforms by @copybara-service[bot] in https://github.com/google/flax/pull/5341
• don't allow Variable mutation in custom_vjp on differentiable arguments by @copybara-service[bot] in https://github.com/google/flax/pull/5342
• support hijax and ref Variables in simple transforms by @copybara-service[bot] in https://github.com/google/flax/pull/5345
• allow in/out_axes when graph_updates=False by @copybara-service[bot] in https://github.com/google/flax/pull/5323
• add transform_metadata transform by @copybara-service[bot] in https://github.com/google/flax/pull/5346
• clean up custom_vjp's graph_updates=False section by @copybara-service[bot] in https://github.com/google/flax/pull/5350
• improve error messages for tree mode duplicates check by @copybara-service[bot] in https://github.com/google/flax/pull/5347
• Introduce manual_type: ManualAxisType parameter on ShapedArray to track varying/unreduced/reduced and remove vma: frozenset parameter. by @copybara-service[bot] in https://github.com/google/flax/pull/5339
• remove aliases from nnx.graph in favor of nnx.compat by @copybara-service[bot] in https://github.com/google/flax/pull/5348
• check aliases on all transform args and simplify apply_variable_updates by @copybara-service[bot] in https://github.com/google/flax/pull/5349
• Add intermediate value captures (extends [#4925]) by @samanklesaria in https://github.com/google/flax/pull/5257
• fix jit_partial lower and compile by @copybara-service[bot] in https://github.com/google/flax/pull/5355
• Added support for data masking in Average, Accuracy and MultiMetric by @vfdev-5 in https://github.com/google/flax/pull/5332
• Do a few more cleanups after pmap_shmap merge by @copybara-service[bot] in https://github.com/google/flax/pull/5358
• update to version 0.12.6 by @copybara-service[bot] in https://github.com/google/flax/pull/5356
• Do a few more cleanups after pmap_shmap merge by @copybara-service[bot] in https://github.com/google/flax/pull/5359
• Do a few more cleanups after the pmap shmap merge deletion. by @copybara-service[bot] in https://github.com/google/flax/pull/5361
• simplify SimpleScan by @copybara-service[bot] in https://github.com/google/flax/pull/5357
• improve nnx.Dict error handling by @copybara-service[bot] in https://github.com/google/flax/pull/5362
• add graph node in prefix checks by @copybara-service[bot] in https://github.com/google/flax/pull/5365
• improve aliasing error msg by @copybara-service[bot] in https://github.com/google/flax/pull/5364
• Generalize out_sharding to work with NamedSharding and Format by @samanklesaria in https://github.com/google/flax/pull/5246
• add nnx.map and nnx.abstract_with_sharding by @copybara-service[bot] in https://github.com/google/flax/pull/5366

Full Changelog: https://github.com/google/flax/compare/v0.12.5...v0.12.6