Skip to content

Pull requests: jax-ml/jax

New pull request New
737 Open 26,633 Closed
737 Open 26,633 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

[Pallas] Turn on --xla_sc_assert_level=san in SC tests.
#37881 opened May 20, 2026 by copybara-service Bot Loading…
Update requirement lock files after 0.10.1 release
#37880 opened May 20, 2026 by kanglant Collaborator Loading…
2
Pass out_sharding through random_bits_p
#37878 opened May 20, 2026 by copybara-service Bot Loading…
@jakevdp
PR for testing
#37876 opened May 20, 2026 by copybara-service Bot Loading…
fix github CI gpu timeouts
#37875 opened May 20, 2026 by copybara-service Bot Loading…
@cgarciae
[SDY] Update PjRt MLIR/HLO conversion APIs to accept a SDY target version during serialization
#37874 opened May 20, 2026 by copybara-service Bot Loading…
[Mosaic GPU] Do not predicate
#37873 opened May 20, 2026 by copybara-service Bot Loading…
@bchetioui
threefry2x32: use unrolled loops in CPU lowering CPU Issues related to the CPU compiler/runtime performance make things lean and fast
#37872 opened May 20, 2026 by jakevdp Collaborator Loading…
@jakevdp
2
Postrelease/0.10.1
#37871 opened May 20, 2026 by kanglant Collaborator Loading…
1
[Mosaic] Support unpack_elementwise to unpack to non-32-bit dtypes.
#37868 opened May 20, 2026 by copybara-service Bot Loading…
Remove the xla::PackIntN and xla::UnpackIntN overloads that are error prone
#37856 opened May 20, 2026 by copybara-service Bot Loading…
Eliminate the use of xla::PackIntN and xla::UnpackIntN overloads that are error prone
#37855 opened May 20, 2026 by copybara-service Bot Loading…
[Pallas TPU] Bugfix: ensure we physicalize output avals during lowering
#37851 opened May 20, 2026 by copybara-service Bot Loading…
@sharadmv
[Pallas] Add tracing_context to Mesh
#37850 opened May 20, 2026 by copybara-service Bot Loading…
@sharadmv
[Pallas] Add physicalize rule for mpmd_map_p
#37848 opened May 20, 2026 by copybara-service Bot Loading…
@sharadmv
Fix JAX fuser physicalize rule argument forwarding bug
#37847 opened May 20, 2026 by copybara-service Bot Loading…
@sharadmv
Linux x86_64 -> Linux aarch64 cross compilation
#37846 opened May 20, 2026 by copybara-service Bot Loading…
@yuriivcs
1
Enable MpmdAsyncTest in OSS tests
#37844 opened May 20, 2026 by copybara-service Bot Loading…
@apaszke
fix debug info crash kokoro:force-run pull ready Ready for copybara import and testing
#37843 opened May 20, 2026 by mattjj Collaborator Loading…
@mattjj
2
Move _get_owner call in stablehlo lowering into the wrapping functions that need it.
#37841 opened May 19, 2026 by copybara-service Bot Loading…
@hawkinsp
[pallas/mgpu] Add a copy_smem_to_gmem scatter operation.
#37839 opened May 19, 2026 by copybara-service Bot Loading…
Reverts c92fe38a323ffd0a9bd26f38f2a0399352e5c769
#37837 opened May 19, 2026 by copybara-service Bot Loading…
Use check_all_safe_to_cast and check_broadcast_shapes for binomial
#37807 opened May 19, 2026 by samanklesaria Contributor • Draft
Add _check_all_safe_to_cast to f sampler
#37806 opened May 19, 2026 by samanklesaria Contributor Loading…
Add dtype and shape checks for dirichlet sampler
#37804 opened May 19, 2026 by samanklesaria Contributor Loading…
6
ProTip! Filter pull requests by the default branch with base:main.