Skip to content

Pull requests: google/flax

New pull request New
268 Open 3,217 Closed
268 Open 3,217 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

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…
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
Handle capture variables in nnx.jit.
#5490 opened Jun 9, 2026 by samanklesaria Collaborator Loading…
6
Add multimodal llm example
#5488 opened Jun 9, 2026 by samanklesaria Collaborator Loading…
3
allow captures in nnx.grad
#5484 opened Jun 4, 2026 by copybara-service Bot Loading…
Add RoPE embeddings pull ready
#5481 opened Jun 2, 2026 by samanklesaria Collaborator Loading…
24
Add docstrings to nnx filterlib filter classes
#5474 opened May 29, 2026 by santihdzs Loading…
3 of 4 tasks
2
ProTip! What’s not been updated in a month: updated:<2026-06-23.