Skip to content

Pull requests: google/flax

New pull request New
258 Open 3,189 Closed
258 Open 3,189 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

Add RoPE embeddings
#5481 opened Jun 2, 2026 by samanklesaria Collaborator Loading…
3
Added Object Detection example with minimal DETR
#5480 opened Jun 2, 2026 by vfdev-5 Collaborator Loading…
1
Fixed typos in docs_nnx/migrating/pytorch_to_jax_flax.rst
#5479 opened Jun 2, 2026 by vfdev-5 Collaborator • Draft
Raise a clear error when pooling window_shape has too many dimensions
#5478 opened Jun 2, 2026 by adityasingh2400 Loading…
1
Fixed gemma example sampler test test_models_from_kaggle
#5477 opened Jun 1, 2026 by vfdev-5 Collaborator Loading…
optimize jit_partial
#5476 opened May 31, 2026 by copybara-service Bot Loading…
nnx.clone: add arrays parameter to also copy underlying buffers
#5475 opened May 29, 2026 by dparikh79 Loading…
4 tasks done
1
1
Add docstrings to nnx filterlib filter classes
#5474 opened May 29, 2026 by santihdzs Loading…
3 of 4 tasks
1
Add sepreformer example
#5472 opened May 26, 2026 by samanklesaria Collaborator Loading…
5
Document spmd_axis_name and transform_metadata in nnx.vmap
#5471 opened May 22, 2026 by adityasingh2400 Loading…
Document None defaults for MultiHeadAttention args
#5470 opened May 21, 2026 by adityasingh2400 Loading…
2 tasks done
1
1
Document LSTMCell carry as the (c, h) tuple it actually is
#5469 opened May 21, 2026 by adityasingh2400 Loading…
3 tasks done
1
1
Fix MGUCell hidden state update
#5467 opened May 21, 2026 by lanarkite99 Loading…
1
2
Migrate away from private jax._src.tree_util._registry
#5464 opened May 20, 2026 by copybara-service Bot Loading…
feat(nnx): add DyT (Dynamic Tanh) normalization module
#5459 opened May 11, 2026 by Ashutosh0x • Draft
1
Add pyrefly type checking: pyrefly.toml config, CI integration
#5458 opened May 11, 2026 by vfdev-5 Collaborator • Draft
3
Added ViT training example pull ready
#5455 opened May 5, 2026 by vfdev-5 Collaborator Loading…
1
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…
Bump the uv group across 10 directories with 5 updates dependencies Pull requests that update a dependency file python:uv Pull requests that update python:uv code
#5432 opened Apr 18, 2026 by dependabot Bot Loading…
Avoid usage of deprecated jax.core APIs.
#5420 opened Apr 15, 2026 by copybara-service Bot Loading…
use State instead of GraphState
#5401 opened Apr 8, 2026 by copybara-service Bot Loading…
use_running_average and deterministic default to None
#5399 opened Apr 7, 2026 by copybara-service Bot Loading…
Consolidate "view" functions with similar naming convention.
#5376 opened Mar 25, 2026 by samanklesaria Collaborator Loading…
3
Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11
#5371 opened Mar 24, 2026 by copybara-service Bot Loading…
1
ProTip! Filter pull requests by the default branch with base:main.