Skip to content

Pull requests: jax-ml/jax

New pull request New
744 Open 23,730 Closed
744 Open 23,730 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

Save CPU version of collective metadata to resources so it can be used during execution.
#34653 opened Jan 27, 2026 by copybara-service bot Loading…
[Pallas/Mosaic GPU] Use self.capture_stdout instead of jtu.capture_stdout in mosaic_gpu_test.
#34652 opened Jan 27, 2026 by copybara-service bot Loading…
@bchetioui
[XLA:GPU] Add CPU version of collective metadata to API.
#34651 opened Jan 27, 2026 by copybara-service bot Loading…
Add a CHLO-based lowering for JAX scan on GPU.
#34650 opened Jan 27, 2026 by copybara-service bot Loading…
Fix: Use ASAN feature instead of config setting
#34647 opened Jan 27, 2026 by copybara-service bot Loading…
@yuriivcs
1
Switch XLA archive download URL from GitHub API to direct archive.
#34646 opened Jan 27, 2026 by copybara-service bot Loading…
@nitins17
[ROCm] Added skip for "test_prim_tridiagonal_solve" tests on ROCm.
#34644 opened Jan 26, 2026 by tsrw2048 Loading…
@emilyfertig
5
[Mosaic GPU][NFC] Pass CompiledKernel by pointer to CachedInit.
#34642 opened Jan 26, 2026 by copybara-service bot Loading…
[WIP] Add support to dynamically set test deps for jax as an external repo
#34641 opened Jan 26, 2026 by alekstheod Loading…
[XLA:MGPU] Port the entire TiledLayout class to C++. CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#34635 opened Jan 26, 2026 by copybara-service bot Loading…
Emit chlo.scan from JAX when targeting GPU and lower to mhlo.scan and only then expand to while loop.
#34631 opened Jan 26, 2026 by copybara-service bot Loading…
fix(scipy): consistently return float and handle nan propagation in rankdata
#34628 opened Jan 26, 2026 by abdulwahabahmedkhanyusufzai Loading…
1
@jakevdp
25
[Doc] Update jnp.argmin/argmax doc for Pallas TPU
#34623 opened Jan 25, 2026 by catswe Loading…
5
Automated Code Change
#34621 opened Jan 24, 2026 by copybara-service bot Loading…
Automated Code Change
#34618 opened Jan 24, 2026 by copybara-service bot Loading…
[ROCm] Skip test_batch_axis_sharding_jvp on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM)
#34615 opened Jan 23, 2026 by AratiGanesh Loading…
1 task
4
Enable miscellaneous tests on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM)
#34603 opened Jan 23, 2026 by phambinhfin Loading…
7
Fix and enable Pallas ops tests on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM)
#34602 opened Jan 23, 2026 by phambinhfin Loading…
Enable neural network tests on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM)
#34601 opened Jan 23, 2026 by phambinhfin Loading…
1
Enable scaled_matmul tests on ROCm AMD GPU Issues pertaining to AMD GPUs (ROCM)
#34600 opened Jan 23, 2026 by phambinhfin Loading…
1
Implement approx_tanh for ROCm using OCML tanh function
#34598 opened Jan 23, 2026 by phambinhfin Loading…
3
Internal-only change.
#34594 opened Jan 23, 2026 by copybara-service bot Loading…
@kanglant
[XLA:MGPU] Port TileTransform to C++.
#34593 opened Jan 23, 2026 by copybara-service bot Loading…
[Pallas/interpreter] Support interpreting non-trivial GPU grids.
#34590 opened Jan 23, 2026 by copybara-service bot Loading…
[Mosaic TPU] Relax offset alignment requirement during mask generation for packed types in SingleRowVRegBounds.
#34585 opened Jan 23, 2026 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.