Skip to content

Pull requests: google/flax

New pull request New
194 Open 2,761 Closed
194 Open 2,761 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

[WIP] Auto mode spmd sharding utilities and docs
#4844 opened Jul 23, 2025 by IvyZX • Draft
[nnx] add Optimizer keep_model param
#4842 opened Jul 23, 2025 by cgarciae • Draft
1
Migrate transforms_test.py from .value to [...]
#4841 opened Jul 23, 2025 by lukeyeh Loading…
Migrate rnglib.py from .value to [...]
#4840 opened Jul 23, 2025 by lukeyeh Loading…
Migrate metrics.py from .value to [...]
#4839 opened Jul 23, 2025 by lukeyeh Loading…
Migrate module_test.py from .value to [...]
#4838 opened Jul 23, 2025 by lukeyeh Loading…
Migrate containers_test.py from .value to [...]
#4837 opened Jul 23, 2025 by lukeyeh Loading…
Migrate conv_test.py from .value to [...]
#4836 opened Jul 23, 2025 by lukeyeh Loading…
Migrate embed_test.py from .value to [...]
#4835 opened Jul 23, 2025 by lukeyeh Loading…
Migrate normalization_test.py from .value to [...]
#4834 opened Jul 23, 2025 by lukeyeh Loading…
Migrate linear_test.py from .value to [...]
#4833 opened Jul 23, 2025 by lukeyeh Loading…
Migrate sampler_test.py from .value to [...]
#4832 opened Jul 23, 2025 by lukeyeh Loading…
Fixed Gemma example using Gemma2 models pull ready
#4830 opened Jul 22, 2025 by vfdev-5 Loading…
1 of 4 tasks
1
3
[nnx] update bridge
#4828 opened Jul 21, 2025 by cgarciae Loading…
[nnx] add flax_pytree_module flag pull ready
#4811 opened Jul 10, 2025 by cgarciae Loading…
Make Sequential() be identity
#4796 opened Jun 28, 2025 by SobhanMP Loading…
3 tasks done
1
6
Add a JAX/Flax key concepts doc Priority: P1 - soon Response within 5 business days. Resolution within 30 days. (Assignee required)
#4795 opened Jun 27, 2025 by IvyZX Loading…
@IvyZX @8bitmp3
4
Use .format in place of .layout.
#4785 opened Jun 19, 2025 by copybara-service bot Loading…
feat: Add overload signatures to module.apply
#4783 opened Jun 18, 2025 by Daraan Loading…
2 tasks done
1
Add BatchNorm layers to CNN in MNIST tutorial for improved training stability
#4773 opened Jun 7, 2025 by sanepunk Loading…
4
Fix typo on 'Scale up on multiple devices' guide
#4765 opened May 30, 2025 by DavidUlloa6310 Loading…
@IvyZX @8bitmp3
2
[nnx] add optax optimizer
#4752 opened May 16, 2025 by cgarciae Loading…
mark shard_map as implemented in transforms guide pull ready
#4738 opened May 5, 2025 by cgarciae Loading…
1
Remove VariableState
#4729 opened Apr 25, 2025 by cgarciae • Draft
1
Fix small typo in performance consideration docs
#4721 opened Apr 20, 2025 by cdelv Loading…
@IvyZX @8bitmp3
2
ProTip! Filter pull requests by the default branch with base:main.