Skip to content

Pull requests: jax-ml/jax

New pull request New
631 Open 26,032 Closed
631 Open 26,032 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

Fix discharge bugs on DMA semaphores
#37041 opened Apr 21, 2026 by copybara-service bot Loading…
@IvyZX
Add dtype to choice function
#37040 opened Apr 21, 2026 by samanklesaria Contributor • Draft
4
Add 'dtype' argument to permutation sampler
#37039 opened Apr 21, 2026 by samanklesaria Contributor Loading…
3
[Pallas MGPU] Add warp-level support to semaphore_signal
#37038 opened Apr 21, 2026 by copybara-service bot Loading…
@Rifur13
[pallas:triton] Compile Pallas Triton kernels to PTX by default
#37034 opened Apr 21, 2026 by copybara-service bot Loading…
@superbobry
[Pallas:MGPU] Add support for remapping SMEM refs within the cluster
#37032 opened Apr 21, 2026 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Unblock int4 WG conversions.
#37027 opened Apr 21, 2026 by copybara-service bot Loading…
[Mosaic GPU] Add support for broadcasting both size-1 and size-0 dimensions.
#37026 opened Apr 21, 2026 by copybara-service bot Loading…
[Mosaic GPU] Allow collapsing unit dimensions when issuing TMA loads.
#37024 opened Apr 21, 2026 by copybara-service bot Loading…
@bchetioui
[Pallas MGPU] Add new MGPU semaphore primitives that take a parameter that specifies the memory scope.
#37018 opened Apr 21, 2026 by copybara-service bot Loading…
@Rifur13
[Pallas MGPU] Use red.add instead of atom.add for semaphore signals
#37017 opened Apr 21, 2026 by copybara-service bot Loading…
@Rifur13
[Mosaic] Add some utility functions for VectorLayout
#37016 opened Apr 21, 2026 by copybara-service bot Loading…
Add jax.scipy.linalg.circulant
#37015 opened Apr 21, 2026 by Saumay Loading…
@jakevdp
8
[Pallas TPU] Add is_remote to dma_wait_p for inserting dummy device ids in the lowering.
#37010 opened Apr 21, 2026 by copybara-service bot Loading…
@sharadmv
[Pallas] Move CoreMemorySpace to Pallas core
#37008 opened Apr 21, 2026 by copybara-service bot Loading…
@sharadmv
Add PRNG key support to colocated_python.
#37004 opened Apr 20, 2026 by copybara-service bot Loading…
Fixes #35493
#37003 opened Apr 20, 2026 by samanklesaria Contributor Loading…
5
[random] Add test scaffolding for out_sharding parameter kokoro:force-run pull ready Ready for copybara import and testing
#37002 opened Apr 20, 2026 by samanklesaria Contributor Loading…
@jakevdp
12
Bump pydata-sphinx-theme from 0.16.1 to 0.17.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#37000 opened Apr 20, 2026 by dependabot bot Loading…
Bump pytest from 8.4.2 to 9.0.3 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#36999 opened Apr 20, 2026 by dependabot bot Loading…
Add out_sharding arguments to jax.random functions
#36993 opened Apr 20, 2026 by samanklesaria Contributor • Draft
@jakevdp
11
Removed a few stale Pyrefly suppressions
#36992 opened Apr 20, 2026 by superbobry Collaborator • Draft
[ROCm] Skip tridiagonal_solve_perturbed code path in eigh unit tests on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM) kokoro:force-run pull ready Ready for copybara import and testing
#36984 opened Apr 20, 2026 by tsrw2048 Contributor Loading…
1
[JAX:GPU] Use is_device_cuda/rocm instead of checking it with test_device_patches in JAX tests
#36982 opened Apr 20, 2026 by copybara-service bot Loading…
[Pallas:GPU] Use is_device_cuda/rocm instead of checking it with test_device_patches in Pallas tests
#36981 opened Apr 20, 2026 by copybara-service bot Loading…
ProTip! What’s not been updated in a month: updated:<2026-03-21.