Skip to content

Pull requests: google/flax

New pull request New
245 Open 3,039 Closed
245 Open 3,039 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

rename graph.py -> graphlib.py
#5290 opened Feb 27, 2026 by copybara-service bot Loading…
[pmap] Remove jax.config.pmap_shmap_merge.
#5289 opened Feb 27, 2026 by copybara-service bot Loading…
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
Add tests for nnx.initializers zeros_init and ones_init
#5276 opened Feb 25, 2026 by us Loading…
2 tasks done
10
Add tests for flax.nnx.nn.dtypes
#5275 opened Feb 25, 2026 by us Loading…
2 tasks done
4
Add PReLU test coverage and move TestPReLUConsistency to activations_test
#5274 opened Feb 25, 2026 by us Loading…
2 tasks done
3
Fixed two typos in the Module & Pytree guide
#5271 opened Feb 23, 2026 by Zirunis Loading…
1 task done
4
Add split method to RngStream
#5270 opened Feb 23, 2026 by samanklesaria Loading…
4
add set_graph_mode
#5269 opened Feb 23, 2026 by copybara-service bot Loading…
feat(nnx): add GQA support to MultiHeadAttention
#5259 opened Feb 18, 2026 by ayulockedin Loading…
3 of 4 tasks
1
10
Add intermediate value captures (extends #4925)
#5257 opened Feb 16, 2026 by samanklesaria Loading…
14
feat(nnx): add out_sharding to recurrent cells
#5255 opened Feb 15, 2026 by aarushisingh04 Loading…
2 of 4 tasks
7
feat(nnx): add preferred_element_type to recurrent cells
#5254 opened Feb 15, 2026 by aarushisingh04 Loading…
2 of 4 tasks
3
feat(nnx): add preferred_element_type to attention module
#5253 opened Feb 15, 2026 by aarushisingh04 Loading…
2 of 4 tasks
5
Add sharding propagation support in nnx.eval_shape (clone of #5111)
#5247 opened Feb 13, 2026 by samanklesaria Loading…
1
5
Generalize out_sharding to work with NamedSharding and Format
#5246 opened Feb 13, 2026 by samanklesaria Loading…
2
Track explicit sharding in Variables
#5245 opened Feb 12, 2026 by samanklesaria Loading…
2
Fix PyTreeNode + Generic losing __parameters__ when Generic is last in bases
#5237 opened Feb 10, 2026 by mohsinm-dev Loading…
2
3
Add 'opt_sharding' option for Variables
#5231 opened Feb 9, 2026 by samanklesaria Loading…
4
ProTip! Add no:assignee to see everything that’s not assigned.