Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
Open
Closed
Improve Pytree flatten/unflatten
#5216 opened 2026-02-02 23:43 by
copybara-service[bot]
Rename sharding_names to out_sharding in NNX Variable metadata
#5215 opened 2026-01-30 23:34 by
copybara-service[bot]
Return updates PyTree from nnx.Optimizer.update
#5214 opened 2026-01-28 23:16 by
anshuljoshi
Add a TAP test for Flax/Keras integration tests.
#5210 opened 2026-01-28 21:51 by
copybara-service[bot]
Add out_sharding argument to Embed layer call
#5205 opened 2026-01-27 17:50 by
samanklesaria
feat(nnx): add object-oriented pooling layers and GlobalAveragePool
#5201 opened 2026-01-26 16:56 by
starryendymion
optax docs test
#5197 opened 2026-01-21 20:32 by
cgarciae
allow nnx standalone import
pull ready
#5195 opened 2026-01-20 22:30 by
cgarciae
This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy).
#5187 opened 2026-01-13 18:30 by
copybara-service[bot]
feat(nnx): add Grouped Query Attention (GQA) support
pull ready
#5180 opened 2026-01-07 23:23 by
ayulockedin
feat(nnx): add preferred_element_type to attention.py and recurrent.py, out_sharding to recurrent.py
#5179 opened 2026-01-07 21:09 by
aarushisingh04
set_mode guide
#5178 opened 2026-01-07 20:23 by
chapman20j
docs(nnx): standardize 'rngs' docstring description
#5173 opened 2026-01-06 18:29 by
ayulockedin
[graphcast] Prepare for `jax_pmap_shmap_merge=True`.
#5168 opened 2026-01-06 06:46 by
copybara-service[bot]
Add out_sharding docstrings to linear layer call methods
#5154 opened 2025-12-25 21:04 by
jackopenn
flatten respect nnx.data
#5140 opened 2025-12-15 17:13 by
cgarciae
Have _graph_flatten respect nnx.data declarations
#5136 opened 2025-12-12 20:19 by
samanklesaria
Add Optimization Cookbook
#5117 opened 2025-11-28 20:02 by
samanklesaria
Added sharding propagation support in nnx.eval_shape
#5111 opened 2025-11-24 15:37 by
vfdev-5
Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put
pull ready
#5101 opened 2025-11-19 17:47 by
vfdev-5
Reenabled examples tests
#5098 opened 2025-11-18 10:58 by
vfdev-5
Refactor nnx.tabulate
#5096 opened 2025-11-18 06:07 by
cgarciae
Fixed dataset path in train tests in examples
#5092 opened 2025-11-17 10:19 by
vfdev-5
Updated flax_gspmd/bridge_guide to use kernel_metadata, make_mesh and sharding_names
#5088 opened 2025-11-13 15:52 by
vfdev-5
Migrate VAE example to Flax NNX
#5077 opened 2025-11-07 20:54 by
sanepunk
add support for hijax Variables in nnx transforms
#5065 opened 2025-11-01 19:19 by
cgarciae
Add split and key methods to RngStream and Rngs
pull ready
#5047 opened 2025-10-23 14:40 by
samanklesaria
Box support
#5029 opened 2025-10-14 23:29 by
cgarciae
New set mode
#5028 opened 2025-10-14 21:12 by
chapman20j
Make Linen and NNX logical rule deduction align
#5003 opened 2025-10-07 01:20 by
IvyZX
Older