Skip to content

Pull requests: google/flax

New pull request New
251 Open 3,110 Closed
251 Open 3,110 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Fixed broken url in the sow warning
#5377 opened Mar 26, 2026 by vfdev-5 Loading…
Consolidate "view" functions with similar naming convention.
#5376 opened Mar 25, 2026 by samanklesaria • Draft
2
[pmap] In-line definitions of jax.device_put_sharded and jax.device_put_replicated.
#5374 opened Mar 25, 2026 by copybara-service bot Loading…
1
Replaced custom PriorityStr with standard tuple sorting in _graph_node_flatten. Python 3.13.11 optimized str comparison bypasses custom subclasses (causing TypeError), and used getattr in _setattr because unflattening order becomes non-deterministic (preventing AttributeError).
#5373 opened Mar 24, 2026 by copybara-service bot Loading…
1
Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11
#5371 opened Mar 24, 2026 by copybara-service bot Loading…
1
Update pytree guide to use nnx.capture
#5370 opened Mar 23, 2026 by samanklesaria Loading…
1
Update the RNG guide with split and without StateAxes
#5363 opened Mar 20, 2026 by samanklesaria Loading…
2
Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone
#5360 opened Mar 20, 2026 by copybara-service bot Loading…
improve scaling up guide
#5354 opened Mar 18, 2026 by cgarciae Loading…
2
Rename SDS(vma: frozenset = ...) to SDS(manual_type: jax.sharding.ManualAxisType = ...)
#5351 opened Mar 18, 2026 by copybara-service bot Loading…
Add 'prefix' argument to temp_flip_flag function
#5343 opened Mar 17, 2026 by samanklesaria Loading…
4
Fix broken error-message formatting in flax.nnx
#5334 opened Mar 16, 2026 by gdrosa Loading…
1 task done
4
Add nnx.prefix
#5325 opened Mar 12, 2026 by samanklesaria Loading…
11
add transforms_tree guide
#5307 opened Mar 4, 2026 by cgarciae Loading…
7
Add iter_module_children() to iterate over module children only.
#5302 opened Mar 4, 2026 by copybara-service bot Loading…
Bump the pip group across 9 directories with 4 updates dependencies Pull requests that update a dependency file python Pull requests that update python code
#5300 opened Mar 4, 2026 by dependabot bot Loading…
change defaults for Dropout and BatchNorm
#5299 opened Mar 3, 2026 by copybara-service bot Loading…
add nnx_basics_tree.ipynb
#5293 opened Mar 2, 2026 by cgarciae Loading…
9
Expand Einsum test coverage with parameterized tests
#5285 opened Feb 26, 2026 by us Loading…
1 of 4 tasks
6
Add tests, docstrings and validation to metrics.py
#5284 opened Feb 26, 2026 by us Loading…
2 of 4 tasks
3
Add examples and fix docstrings in nnx.nn.dtypes
#5283 opened Feb 26, 2026 by us Loading…
2 of 4 tasks
3
Expand Dropout test coverage and migrate to absltest
#5281 opened Feb 26, 2026 by us Loading…
1 of 4 tasks
3
Expand LoRA test coverage with dtype, gradient, and edge case tests
#5280 opened Feb 26, 2026 by us Loading…
1 of 4 tasks
4
Expand TestLinenConsistency parameterized.product in attention_test
#5279 opened Feb 26, 2026 by us Loading…
2 of 4 tasks
1
Expand Embed test coverage with functional, attend, and dtype tests
#5278 opened Feb 26, 2026 by us Loading…
1 of 4 tasks
3
ProTip! Updated in the last three days: updated:>2026-03-23.