Thanks to visit codestin.com
Credit goes to Github.com

Skip to content

Pull requests: google/flax

New pull request New
232 Open 3,022 Closed
232 Open 3,022 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

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 tree-mode support to nnx.{cond,switch,eval_shape,checkify}
#5252 opened Feb 14, 2026 by copybara-service bot Loading…
Add tree-mode arg error handling
#5251 opened Feb 14, 2026 by copybara-service bot Loading…
Improve unregistered data detection
#5249 opened Feb 13, 2026 by copybara-service bot Loading…
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
[flax:benchmarks] Fix flax benchmarks.
#5241 opened Feb 12, 2026 by copybara-service bot Loading…
Fix PyTreeNode + Generic losing __parameters__ when Generic is last in bases
#5237 opened Feb 10, 2026 by mohsinm-dev Loading…
1
3
Add 'opt_sharding' option for Variables
#5231 opened Feb 9, 2026 by samanklesaria Loading…
4
fix: correct typo 'seperate' to 'separate'
#5229 opened Feb 9, 2026 by thecaptain789 Loading…
4
Removed redundant code checking wrt in nnx.Optimizer
#5226 opened Feb 6, 2026 by vfdev-5 Loading…
1
enable mlp test
#5224 opened Feb 6, 2026 by copybara-service bot Loading…
Remove nnx.List
#5223 opened Feb 4, 2026 by samanklesaria Loading…
25
Return updates PyTree from nnx.Optimizer.update
#5214 opened Jan 28, 2026 by anshuljoshi Loading…
3 tasks done
1
4
Add a TAP test for Flax/Keras integration tests.
#5210 opened Jan 28, 2026 by copybara-service bot Loading…
feat(nnx): add object-oriented pooling layers and GlobalAveragePool
#5201 opened Jan 26, 2026 by divye-joshi Loading…
1
3
optax docs test
#5197 opened Jan 21, 2026 by cgarciae • Draft
allow nnx standalone import pull ready
#5195 opened Jan 20, 2026 by cgarciae Loading…
1
This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy).
#5187 opened Jan 13, 2026 by copybara-service bot Loading…
feat(nnx): add preferred_element_type to attention.py and recurrent.py, out_sharding to recurrent.py
#5179 opened Jan 7, 2026 by aarushisingh04 Loading…
2 of 4 tasks
20
set_mode guide
#5178 opened Jan 7, 2026 by chapman20j Loading…
1 of 4 tasks
1
[graphcast] Prepare for jax_pmap_shmap_merge=True.
#5168 opened Jan 6, 2026 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.