Skip to content

Pull requests: google/flax

New pull request New
263 Open 3,209 Closed
263 Open 3,209 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: 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 Loading…
4 tasks done
1
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 • Draft
fix(nnx): raise ValueError instead of assert for q/k/v shape checks in attention
#5514 opened Jun 25, 2026 by Sumu004 Loading…
3 tasks done
1
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
Validate dropout_rate in dot_product_attention
#5503 opened Jun 16, 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
Add vmap reset() regression test for nnx.metrics (#5483)
#5491 opened Jun 11, 2026 by chenkuanliao Loading…
4 tasks done
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…
1
allow captures in nnx.grad
#5484 opened Jun 4, 2026 by copybara-service Bot Loading…
feat(nnx): add arrays=True to nnx.clone for independent buffer copies
#5482 opened Jun 4, 2026 by Sumu004 Loading…
1
Add RoPE embeddings
#5481 opened Jun 2, 2026 by samanklesaria Collaborator Loading…
22
nnx.clone: add arrays parameter to also copy underlying buffers
#5475 opened May 29, 2026 by dparikh79 Loading…
4 tasks done
1
@vfdev-5
4
Add docstrings to nnx filterlib filter classes
#5474 opened May 29, 2026 by santihdzs Loading…
3 of 4 tasks
2
Add sepreformer example
#5472 opened May 26, 2026 by samanklesaria Collaborator Loading…
6
Fix MGUCell hidden state update
#5467 opened May 21, 2026 by lanarkite99 Loading…
1
3
Migrate away from private jax._src.tree_util._registry
#5464 opened May 20, 2026 by copybara-service Bot Loading…
Add pyrefly type checking: pyrefly.toml config, CI integration
#5458 opened May 11, 2026 by vfdev-5 Collaborator • Draft
3
Bump the pip group across 10 directories with 4 updates dependencies Pull requests that update a dependency file python Pull requests that update python code
#5444 opened Apr 24, 2026 by dependabot Bot Loading…
migrate to pyrefly
#5437 opened Apr 22, 2026 by copybara-service Bot Loading…
ProTip! Adding no:label will show everything without a label.