Skip to content

Pull requests: google/flax

New pull request New
223 Open 2,952 Closed
223 Open 2,952 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

[pmap] Avoid degraded performance under the new jax.pmap.
#5152 opened Dec 21, 2025 by copybara-service bot Loading…
Flax 0.12.2 release
#5150 opened Dec 18, 2025 by copybara-service bot Loading…
1
Fixed typos in the docstrings using antigravity
#5145 opened Dec 18, 2025 by vfdev-5 Loading…
Typo fix in _cached_partial method
#5142 opened Dec 18, 2025 by vfdev-5 • Draft
Handle pure bodies in nnx.fori_loop
#5141 opened Dec 17, 2025 by samanklesaria Loading…
1
3
flatten respect nnx.data
#5140 opened Dec 15, 2025 by cgarciae Loading…
fix general guides landing page
#5139 opened Dec 14, 2025 by ivrolan Loading…
2 of 4 tasks
1
2
Have _graph_flatten respect nnx.data declarations
#5136 opened Dec 12, 2025 by samanklesaria • Draft
1
fix kw_only_dataclasses for python 3.14 (part 2)
#5135 opened Dec 11, 2025 by copybara-service bot Loading…
Make nnx.pop remove sown attributes.
#5133 opened Dec 9, 2025 by samanklesaria Loading…
1
3
Support multiple None and UNCONSTRAINED when resolving logical rules
#5129 opened Dec 8, 2025 by copybara-service bot Loading…
Add Optimization Cookbook
#5117 opened Nov 28, 2025 by samanklesaria • Draft
1
improve hijax guide
#5115 opened Nov 25, 2025 by cgarciae 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
Rename sharding_names to sharding_metadata
#5089 opened Nov 13, 2025 by samanklesaria Loading…
1
Updated flax_gspmd/bridge_guide to use kernel_metadata, make_mesh and sharding_names
#5088 opened Nov 13, 2025 by vfdev-5 Loading…
1
Support for python 3.14 pull ready
#5087 opened Nov 13, 2025 by vfdev-5 Loading…
1
12
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
Use nnx split during tabulate pull ready
#5069 opened Nov 3, 2025 by samanklesaria Loading…
16
ProTip! Add no:assignee to see everything that’s not assigned.