Skip to content

Pull requests: google/flax

New pull request New
217 Open 2,948 Closed
217 Open 2,948 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

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…
Copybara import of the project:
#5134 opened Dec 9, 2025 by copybara-service bot Loading…
1
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
add support for hijax Variables in nnx transforms
#5065 opened Nov 1, 2025 by cgarciae Loading…
Update mnist example to use NNX (clone of #5002) pull ready
#5064 opened Oct 31, 2025 by samanklesaria Loading…
1
3
Add split and key methods to RngStream and Rngs pull ready
#5047 opened Oct 23, 2025 by samanklesaria Loading…
1
5
Box support
#5029 opened Oct 14, 2025 by cgarciae • Draft
1
New set mode
#5028 opened Oct 14, 2025 by chapman20j Loading…
1 task
Make Linen and NNX logical rule deduction align
#5003 opened Oct 7, 2025 by IvyZX Loading…
ProTip! Add no:assignee to see everything that’s not assigned.