Skip to content

Pull requests: jax-ml/jax

New pull request New
649 Open 23,240 Closed
649 Open 23,240 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

[test] fix toeplitz test under scipy 1.17 kokoro:force-run pull ready Ready for copybara import and testing
#33956 opened Dec 15, 2025 by jakevdp Loading…
1
@jakevdp
3
Disable a JAX test under TSAN on CPU due to timeouts.
#33955 opened Dec 15, 2025 by copybara-service bot Loading…
@belitskiy
[PJRT] Simplify and broaden the interface of TransposePlan.
#33954 opened Dec 15, 2025 by copybara-service bot Loading…
@hawkinsp
Enable xla_gpu_experimental_scaled_dot_with_triton by default.
#33952 opened Dec 15, 2025 by copybara-service bot Loading…
Suppress "healthcheck too slow" for OpsTest.test_select_n under ASAN.
#33951 opened Dec 15, 2025 by copybara-service bot Loading…
@belitskiy
Automated Code Change
#33946 opened Dec 15, 2025 by copybara-service bot Loading…
Remove more deprecated BUILD aliases.
#33945 opened Dec 14, 2025 by copybara-service bot Loading…
@hawkinsp
Reverts a99f76f6a6556223d71a2a63cd03019634085cc5
#33940 opened Dec 13, 2025 by copybara-service bot Loading…
Fix vjp type annotation
#33938 opened Dec 13, 2025 by copybara-service bot Loading…
@yashk2810
Update rules_ml_toolchain version.
#33935 opened Dec 13, 2025 by copybara-service bot Loading…
@ybaturina
1
[PjRt-IFRT] Create ifrt::PjRtExecutable only from ifrt::PjRtCompiler and CompileOnlyIfrtCompiler
#33933 opened Dec 13, 2025 by copybara-service bot Loading…
Test NVSHMEM fix
#33932 opened Dec 13, 2025 by copybara-service bot Loading…
@ybaturina
1
Deprecate a number of rarely-used jax.core APIs kokoro:force-run pull ready Ready for copybara import and testing
#33931 opened Dec 13, 2025 by jakevdp Loading…
@jakevdp
3
[vjp3] remove ad.backward_pass and non-fancy HOP transposes kokoro:force-run pull ready Ready for copybara import and testing
#33929 opened Dec 12, 2025 by mattjj Loading…
@mattjj
7
[Mosaic:TPU] Clean up tpu.memref_slice verifier
#33927 opened Dec 12, 2025 by copybara-service bot Loading…
Add a thread guard config option.
#33926 opened Dec 12, 2025 by copybara-service bot Loading…
@emilyfertig
[easy] [XLA] JAX should pass tiling as well for it to be respected
#33925 opened Dec 12, 2025 by copybara-service bot Loading…
[Mosaic:TPU] tpu.memref_slice with strided layouts always has a dynamic offset, except for rank == 0
#33923 opened Dec 12, 2025 by copybara-service bot Loading…
[sparse] make caveat at top more prominent documentation kokoro:force-run pull ready Ready for copybara import and testing
#33922 opened Dec 12, 2025 by jakevdp Loading…
@jakevdp
3
[Mosaic:TPU] Clean up tpu.memref_slice verifier
#33917 opened Dec 12, 2025 by copybara-service bot Loading…
Add name field to K8sCluster
#33916 opened Dec 12, 2025 by JKSPL Loading…
1
[Pallas][Mosaic GPU] Fix race condition in dynamic_scheduling_loop
#33914 opened Dec 12, 2025 by copybara-service bot Loading…
@justinjfu
[Pallas][Mosaic GPU] Fix race condition in dynamic_scheduling_loop.
#33913 opened Dec 12, 2025 by copybara-service bot Loading…
@justinjfu
[pxla] Deprecate jax.interpreters.pxla symbols.
#33911 opened Dec 12, 2025 by copybara-service bot Loading…
@danielsuo
test for hunting cuda error
#33909 opened Dec 12, 2025 by copybara-service bot Loading…
ProTip! What’s not been updated in a month: updated:<2025-11-15.