| Name | Modified | Size | Downloads / Week |
|---|---|---|---|
| Parent folder | |||
| 0.12.6 source code.tar.gz | 2026-03-20 | 5.1 MB | |
| 0.12.6 source code.zip | 2026-03-20 | 5.5 MB | |
| README.md | 2026-03-20 | 5.8 kB | |
| Totals: 3 Items | 10.5 MB | 0 | |
What's Changed
- Annotate
wrapped_fn'sselfargument indecorator_lift_transformfor 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: ManualAxisTypeparameter on ShapedArray to track varying/unreduced/reduced and removevma: frozensetparameter. 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