Skip to content

Pull requests: google/flax

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

docs: Document allow_duplicates argument of nnx.to_arrays.
#5083 opened Nov 8, 2025 by dan-zheng Loading…
1 of 4 tasks
1
fix main pull ready
#5081 opened Nov 7, 2025 by cgarciae Loading…
add promote_dtype to all standard layers
#5080 opened Nov 7, 2025 by cgarciae Loading…
Add nnx.use_eager_sharding context manager
#5079 opened Nov 7, 2025 by samanklesaria Loading…
2
Migrate VAE example to Flax NNX with JIT optimization - 9.2x speedup
#5077 opened Nov 7, 2025 by sanepunk Loading…
5 tasks done
17
Added missing LayerNorm test case into TestLayersSameGraph pull ready
#5076 opened Nov 7, 2025 by vfdev-5 Loading…
Added kernel_metadata/bias_metadata args to nnx layers
#5074 opened Nov 6, 2025 by vfdev-5 Loading…
5 tasks done
5
Fixed nnx.scan with carry as pytree and sow
#5073 opened Nov 4, 2025 by vfdev-5 Loading…
1
6
fix List and improve Pytree
#5072 opened Nov 4, 2025 by cgarciae Loading…
Adding set_mode_info pull ready
#5071 opened Nov 4, 2025 by chapman20j Loading…
2 tasks
support explicit sharding in eager sharding
#5070 opened Nov 3, 2025 by cgarciae Loading…
2
Fallback to non-jit version of methods in tabulate if necessary pull ready
#5069 opened Nov 3, 2025 by samanklesaria Loading…
15
add nnx.dataclass
#5066 opened Nov 1, 2025 by cgarciae Loading…
9
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
11
Fix bound method auto-unbinding for NNX transforms pull ready
#5055 opened Oct 28, 2025 by mohsinm-dev Loading…
16
deprecate Variable.value
#5052 opened Oct 24, 2025 by cgarciae Loading…
1
Add split and key methods to RngStream and Rngs pull ready
#5047 opened Oct 23, 2025 by samanklesaria Loading…
1
3
Box support
#5029 opened Oct 14, 2025 by cgarciae • Draft
1
New set mode
#5028 opened Oct 14, 2025 by chapman20j Loading…
1 task
[flax:examples:vae] Small linter fixes.
#5014 opened Oct 9, 2025 by copybara-service bot Loading…
[flax:examples:gemma] Fixing linter errors.
#5013 opened Oct 9, 2025 by copybara-service bot Loading…
[flax:examples:wmt] Small linter fixes.
#5012 opened Oct 9, 2025 by copybara-service bot Loading…
Make Linen and NNX logical rule deduction align
#5003 opened Oct 7, 2025 by IvyZX Loading…
ProTip! Follow long discussions with comments:>50.