Skip to content

Pull requests: google/flax

New pull request New
221 Open 2,984 Closed
221 Open 2,984 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

Remove ref cycles introduced by self-calling nested functions.
#5193 opened Jan 17, 2026 by copybara-service bot Loading…
fix(nnx): ensure PathContains('') matches all paths
#5192 opened Jan 16, 2026 by ayulockedin Loading…
3 of 4 tasks
1
1
Add HijaxTransformCoverageTest
#5190 opened Jan 15, 2026 by cgarciae Loading…
Shard split rngs in lifted vmap if spmd_axis_name is given and applied to vmapped axis. Without this, jax.vmap gives an error: ValueError: Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: None vs batch due to split_rngs being replicated.
#5189 opened Jan 15, 2026 by copybara-service bot Loading…
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 Grouped Query Attention (GQA) support
#5180 opened Jan 7, 2026 by ayulockedin Loading…
3 of 4 tasks
1
@samanklesaria
22
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
10
set_mode guide
#5178 opened Jan 7, 2026 by chapman20j Loading…
1 of 4 tasks
1
docs(nnx): standardize 'rngs' docstring description
#5173 opened Jan 6, 2026 by ayulockedin Loading…
1 of 4 tasks
1
Add _graph_node_set_key method for List class pull ready
#5171 opened Jan 6, 2026 by samanklesaria Loading…
1
[graphcast] Prepare for jax_pmap_shmap_merge=True.
#5168 opened Jan 6, 2026 by copybara-service bot Loading…
refactor: add missing type hints to flax/core/scope.py
#5167 opened Jan 5, 2026 by ayulockedin Loading…
1 of 4 tasks
1
Add out_sharding docstrings to linear layer call methods
#5154 opened Dec 25, 2025 by jackopenn Loading…
2 of 4 tasks
flatten respect nnx.data
#5140 opened Dec 15, 2025 by cgarciae Loading…
Have _graph_flatten respect nnx.data declarations
#5136 opened Dec 12, 2025 by samanklesaria • Draft
1
Add Optimization Cookbook
#5117 opened Nov 28, 2025 by samanklesaria Loading…
1
Added sharding propagation support in nnx.eval_shape
#5111 opened Nov 24, 2025 by vfdev-5 • Draft
1
Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put pull ready
#5101 opened Nov 19, 2025 by vfdev-5 Loading…
Reenabled examples tests
#5098 opened Nov 18, 2025 by vfdev-5 • Draft
1
Refactor nnx.tabulate
#5096 opened Nov 18, 2025 by cgarciae Loading…
Fixed dataset path in train tests in examples
#5092 opened Nov 17, 2025 by vfdev-5 • Draft
Updated flax_gspmd/bridge_guide to use kernel_metadata, make_mesh and sharding_names
#5088 opened Nov 13, 2025 by vfdev-5 • Draft
2
Updated pyproject.toml
#5086 opened Nov 13, 2025 by vfdev-5 Loading…
2
Migrate VAE example to Flax NNX
#5077 opened Nov 7, 2025 by sanepunk Loading…
5 tasks done
23
fix List and improve Pytree
#5072 opened Nov 4, 2025 by cgarciae Loading…
2
ProTip! What’s not been updated in a month: updated:<2025-12-17.