Skip to content

Pull requests: google/flax

New pull request New
272 Open 3,221 Closed
272 Open 3,221 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(nnx): fill missing docstrings in variables/rnglib, fix helpers g…
#5545 opened Aug 5, 2026 by greekera1000 Loading…
3 of 4 tasks
docs(nnx): fix training API reference inconsistencies
#5544 opened Aug 4, 2026 by greekera1000 Loading…
3 of 4 tasks
2
Make Flax thread-safe for multi-threaded tracing.
#5543 opened Aug 4, 2026 by copybara-service Bot Loading…
Enable full NumPy indexing compliance for jax.Ref
#5542 opened Aug 1, 2026 by copybara-service Bot Loading…
Fix vmap axes for input_positions in dot_product_attention_with_rope
#5537 opened Jul 26, 2026 by TheSaiEaranti Loading…
3 tasks done
Fix broken links
#5533 opened Jul 22, 2026 by copybara-service Bot Loading…
2
Exclude haiku doctests from CI
#5532 opened Jul 22, 2026 by samanklesaria Collaborator Loading…
2
Handle captures in jax transforms with hijax
#5531 opened Jul 21, 2026 by samanklesaria Collaborator Loading…
Preserve RngStream counters for passthrough streams in Rngs.split/broadcast
#5529 opened Jul 17, 2026 by sohumt123 Loading…
3 of 4 tasks
1
Fix nnx.Dropout crash when broadcast_dims includes a sharded axis
#5528 opened Jul 17, 2026 by sohumt123 Loading…
3 of 4 tasks
1
Handle captures in nnx.vmap
#5527 opened Jul 7, 2026 by samanklesaria Collaborator Loading…
Fix tracing cache misses in nnx.cond and nnx.switch (#5512)
#5525 opened Jul 6, 2026 by kapellirohith Loading…
1
Add pickle/cloudpickle round-trip tests for Linen modules (#1481)
#5524 opened Jul 6, 2026 by kapellirohith Loading…
2 of 4 tasks
1
1
Speed up NNX graph flattening by avoiding unnecessary array-ref checks
#5523 opened Jul 6, 2026 by KookiesNKareem Loading…
1 of 4 tasks
Handle captures grad
#5522 opened Jul 2, 2026 by samanklesaria Collaborator Loading…
Update NNX docs for deprecated Variable.value access
#5521 opened Jul 2, 2026 by sumitdotml Loading…
2 of 4 tasks
2
docs: clarify shape dynamics and add example for make_causal_mask
#5519 opened Jun 27, 2026 by SakshieP Loading…
3 of 4 tasks
1
1
Fix nnx.cond tracing cache misses by caching SimpleCondFn wrappers
#5518 opened Jun 27, 2026 by chenkuanliao Contributor Loading…
4 tasks done
1
27
cache cond/switch branch wrappers to avoid tracing cache misses
#5517 opened Jun 27, 2026 by mohsinm-dev Contributor Loading…
1
Handle captures in nnx.scan
#5515 opened Jun 25, 2026 by samanklesaria Collaborator Loading…
Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused LayerNorm kernel, custom VJP for reverse-mode autodiff, and numerical correctness benchmark script.
#5513 opened Jun 22, 2026 by copybara-service Bot Loading…
1
fix(attention): validate dropout_rate range and require dropout_rng w…
#5510 opened Jun 21, 2026 by Vitaliy-Pikalo Loading…
3 tasks done
1
1
default partition metadata in vmap/scan/pmap
#5508 opened Jun 18, 2026 by mohsinm-dev Contributor Loading…
1
1
add nnx:v1
#5502 opened Jun 16, 2026 by copybara-service Bot Loading…
fix cached_partial for raw array attributes
#5501 opened Jun 13, 2026 by discobot Loading…
4 tasks done
1
1
ProTip! Add no:assignee to see everything that’s not assigned.