Skip to content

Pull requests: google/flax

New pull request New
256 Open 3,155 Closed
256 Open 3,155 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

Added out_sharding and qkv_sharding to MHA call method
#5434 opened Apr 21, 2026 by vfdev-5 Collaborator • Draft
4 tasks
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…
Add machine translation example
#5431 opened Apr 17, 2026 by samanklesaria Collaborator Loading…
1
Updated gemma docs
#5427 opened Apr 17, 2026 by vfdev-5 Collaborator • Draft
1
Update mnist tutorial to use tensorboard
#5425 opened Apr 17, 2026 by samanklesaria Collaborator Loading…
6
Reorganize examples
#5424 opened Apr 16, 2026 by samanklesaria Collaborator Loading…
5
Fix doubled 'the the' typos across multiple files pull ready
#5422 opened Apr 16, 2026 by SAY-5 Loading…
1
Avoid usage of deprecated jax.core APIs.
#5420 opened Apr 15, 2026 by copybara-service Bot Loading…
Update view guide to mention with_ and as_ prefixes
#5416 opened Apr 14, 2026 by samanklesaria Collaborator Loading…
1
Fix intermediates for Gemma example pull ready
#5409 opened Apr 13, 2026 by vfdev-5 Collaborator Loading…
3
[Docs] migration guide "From PyTorch to JAX and Flax"
#5408 opened Apr 13, 2026 by vfdev-5 Collaborator Loading…
Port Jax AI stack minigpt example
#5405 opened Apr 10, 2026 by samanklesaria Collaborator Loading…
12
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…
Rename vars_as to with_vars
#5384 opened Mar 30, 2026 by samanklesaria Collaborator Loading…
2
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
Update the RNG guide with split and without StateAxes
#5363 opened Mar 20, 2026 by samanklesaria Collaborator Loading…
3
Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone
#5360 opened Mar 20, 2026 by copybara-service Bot Loading…
improve scaling up guide
#5354 opened Mar 18, 2026 by cgarciae Collaborator Loading…
2
Rename SDS(vma: frozenset = ...) to SDS(manual_type: jax.sharding.ManualAxisType = ...)
#5351 opened Mar 18, 2026 by copybara-service Bot Loading…
Add 'prefix' argument to temp_flip_flag function
#5343 opened Mar 17, 2026 by samanklesaria Collaborator Loading…
4
Fix broken error-message formatting in flax.nnx
#5334 opened Mar 16, 2026 by gdrosa Loading…
1 task done
4
add transforms_tree guide
#5307 opened Mar 4, 2026 by cgarciae Collaborator Loading…
7
Bump the pip group across 9 directories with 4 updates dependencies Pull requests that update a dependency file python Pull requests that update python code
#5300 opened Mar 4, 2026 by dependabot Bot Loading…
ProTip! Type g i on any issue or pull request to go back to the issue listing page.