Skip to content

Pull requests: google/flax

New pull request New
209 Open 2,851 Closed
209 Open 2,851 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

propagate eager sharding
#4983 opened Sep 25, 2025 by cgarciae Loading…
[nnx] preserve the function's type information in jit pull ready
#4981 opened Sep 25, 2025 by cgarciae Loading…
3
Move ArrayRef creation to the end of Variable creation pull ready
#4980 opened Sep 24, 2025 by IvyZX Loading…
2
chore: more specific migration guide description
#4977 opened Sep 24, 2025 by emergenz Loading…
1 task done
Update tutorial examples to thread explicit RNGs (WIP)
#4975 opened Sep 23, 2025 by samanklesaria • Draft
1
1
ignore uv.lock
#4974 opened Sep 23, 2025 by copybara-service bot Loading…
add Variable.set_metadata pull ready
#4968 opened Sep 19, 2025 by cgarciae Loading…
Move iter_children and iter_modules to functions
#4961 opened Sep 19, 2025 by samanklesaria Loading…
1
10
Hijax Variable
#4960 opened Sep 18, 2025 by cgarciae Loading…
1
Add compute_flops and compute_vjp_flops options to nnx.tabulate pull ready
#4948 opened Sep 15, 2025 by samanklesaria Loading…
1
45
Simplify the effort to create sharded Flax NNX models.
#4941 opened Sep 12, 2025 by copybara-service bot Loading…
1
Clarify *Norm layer docstrings: axis_index_groups is unused under SPMD jit.
#4940 opened Sep 12, 2025 by copybara-service bot Loading…
Add a customizable conv_transpose. pull ready
#4938 opened Sep 11, 2025 by copybara-service bot Loading…
1
Don't pass linen_meta_type argument when creating AxisMetadata subclasses.
#4933 opened Sep 9, 2025 by copybara-service bot Loading…
1
Improve error message with class name.
#4932 opened Sep 9, 2025 by copybara-service bot Loading…
1
Set mode
#4927 opened Sep 5, 2025 by cgarciae Loading…
adds recursive_map
#4926 opened Sep 5, 2025 by copybara-service bot Loading…
Add intermediate value captures based off hijax.Box
#4925 opened Sep 4, 2025 by IvyZX Loading…
6
Updates in gemma example: grain dataflow, simplified sharding rules
#4913 opened Aug 27, 2025 by vfdev-5 • Draft
1
[jax:benchmarks] Add tracing/lowering benchmarks for a few flax examples.
#4911 opened Aug 26, 2025 by copybara-service bot Loading…
Refactor ArrayCheckpointHandler to use BasePyTreeCheckpointHandler.
#4902 opened Aug 22, 2025 by copybara-service bot Loading…
1
Fix nnx.tabulate crash with empty dict/None values (fixes #4889)
#4891 opened Aug 14, 2025 by mohsinm-dev Loading…
2 of 4 tasks
1
6
Add python 3.13 CI tests
#4888 opened Aug 13, 2025 by IvyZX Loading…
fix_checkpoint helper
#4875 opened Aug 6, 2025 by rcrowe-google Loading…
2 tasks
Fixing formatting
#4874 opened Aug 5, 2025 by rcrowe-google Loading…
4 tasks
2
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.