Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
Open
Closed
Reduce warning message of about GCS usage
#5181 opened 2026-01-08 17:21 by
wdhongtw
feat(nnx): add Grouped Query Attention (GQA) support
#5180 opened 2026-01-07 23:23 by
ayulockedin
feat(nnx): add preferred_element_type and out_sharding support to Attention and RNN modules
#5179 opened 2026-01-07 21:09 by
aarushisingh04
set_mode guide
#5178 opened 2026-01-07 20:23 by
chapman20j
docs(nnx): add missing functional args to Conv and LinearGeneral
#5174 opened 2026-01-06 18:51 by
ayulockedin
docs(nnx): standardize 'rngs' docstring description
#5173 opened 2026-01-06 18:29 by
ayulockedin
Add _graph_node_set_key method for List class
#5171 opened 2026-01-06 16:34 by
samanklesaria
fix: remove superfluous `jax.Ref` standalone statement
#5169 opened 2026-01-06 15:16 by
pfackeldey
[graphcast] Prepare for `jax_pmap_shmap_merge=True`.
#5168 opened 2026-01-06 06:46 by
copybara-service[bot]
refactor: add missing type hints to flax/core/scope.py
#5167 opened 2026-01-05 23:55 by
ayulockedin
Docs: Fix typo and clarify introduction in Functional API section
#5157 opened 2025-12-26 10:23 by
Moriyuki-S
Add out_sharding docstrings to linear layer call methods
#5154 opened 2025-12-25 21:04 by
jackopenn
Fixed typos in the docstrings using antigravity
#5145 opened 2025-12-18 14:05 by
vfdev-5
Typo fix in _cached_partial method
#5142 opened 2025-12-18 08:55 by
vfdev-5
Handle pure bodies in nnx.fori_loop
#5141 opened 2025-12-17 19:29 by
samanklesaria
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
Updated pyproject.toml
#5086 opened 2025-11-13 10:45 by
vfdev-5
Migrate VAE example to Flax NNX
#5077 opened 2025-11-07 20:54 by
sanepunk
fix List and improve Pytree
#5072 opened 2025-11-04 03:10 by
cgarciae
Use nnx split during `tabulate`
pull ready
#5069 opened 2025-11-03 21:50 by
samanklesaria
add support for hijax Variables in nnx transforms
#5065 opened 2025-11-01 19:19 by
cgarciae
Update mnist example to use NNX
pull ready
#5064 opened 2025-10-31 17:53 by
samanklesaria
Older