google/flax

Pull Requests Commits
Open Closed
fix eval_shape's _to_variable
#5321 opened 2026-03-09 17:41 by copybara-service[bot]
add graph_updates to vmap
#5320 opened 2026-03-09 04:30 by copybara-service[bot]
add graph_updates argument to jit
#5317 opened 2026-03-06 18:40 by copybara-service[bot]
add tree-mode-nnx FLIP
#5310 opened 2026-03-04 22:34 by copybara-service[bot]
add transforms_tree guide
#5307 opened 2026-03-04 18:22 by cgarciae
Bump the pip group across 9 directories with 4 updates dependencies python
#5300 opened 2026-03-04 00:07 by dependabot[bot]
add nnx_basics_tree.ipynb
#5293 opened 2026-03-02 19:16 by cgarciae
Expand Einsum test coverage with parameterized tests
#5285 opened 2026-02-26 21:35 by us
Add tests, docstrings and validation to metrics.py
#5284 opened 2026-02-26 21:33 by us
Add examples and fix docstrings in nnx.nn.dtypes
#5283 opened 2026-02-26 21:11 by us
Expand Dropout test coverage and migrate to absltest
#5281 opened 2026-02-26 20:32 by us
Add tests for flax.nnx.nn.dtypes
#5275 opened 2026-02-25 22:46 by us
Fixed two typos in the Module & Pytree guide
#5271 opened 2026-02-23 23:58 by Zirunis
Add split method to RngStream
#5270 opened 2026-02-23 15:53 by samanklesaria
Track explicit sharding in Variables
#5245 opened 2026-02-12 23:02 by samanklesaria
Add 'opt_sharding' option for Variables
#5231 opened 2026-02-09 20:47 by samanklesaria
Terms Of Use Privacy Policy Refund Policy Impressum