Skip to content

Pull requests: google/flax

New pull request New
218 Open 2,938 Closed
218 Open 2,938 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

Avoid passing concrete argument to jax.remat
#5121 opened Dec 2, 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 • Draft
1
Added sharding propagation support in nnx.eval_shape
#5111 opened Nov 24, 2025 by vfdev-5 • Draft
1
Add out_sharding argument to call methods for layers with jax calls that support it pull ready
#5102 opened Nov 19, 2025 by samanklesaria Loading…
2
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…
Added is_causal mask argument to flax.nnx.dot_product_attention
#5093 opened Nov 17, 2025 by ibbyml Loading…
3 tasks done
11
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
1
Pooling operation
#5057 opened Oct 28, 2025 by jorisSchaller Loading…
1 task done
1
13
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…
chore: more specific migration guide description
#4977 opened Sep 24, 2025 by emergenz Loading…
1 task done
ProTip! What’s not been updated in a month: updated:<2025-11-03.