Skip to content

Pull requests: jax-ml/jax

New pull request New
691 Open 21,314 Closed
691 Open 21,314 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

Disallow passing an Array as a dtype argument.
#31797 opened Sep 13, 2025 by copybara-service bot Loading…
@hawkinsp
Remove a canonicalize_value call from partial_eval.py.
#31796 opened Sep 13, 2025 by copybara-service bot Loading…
@hawkinsp
Automated Code Change
#31794 opened Sep 13, 2025 by copybara-service bot Loading…
WIP: Open-source multiprocess test runner.
#31791 opened Sep 13, 2025 by copybara-service bot Loading…
@emilyfertig
[hijax] support refs to hijax types
#31790 opened Sep 13, 2025 by mattjj Loading…
Run oldest supported numpy test with OMP_NUM_THREADS=8
#31789 opened Sep 12, 2025 by copybara-service bot Loading…
@quoctruong
Tag a number of shard_map tests as slow.
#31788 opened Sep 12, 2025 by copybara-service bot Loading…
@hawkinsp
Add x64 checking/warnings when abstractifying ShapeDtypeStruct and its relatives.
#31786 opened Sep 12, 2025 by copybara-service bot Loading…
@hawkinsp
Simplify check_and_canonicalize_user_dtype.
#31785 opened Sep 12, 2025 by copybara-service bot Loading…
@hawkinsp
[multihost_utils] Add PRNGKeyArray support to host_local_array_to_global_array and global_array_to_host_local_array.
#31784 opened Sep 12, 2025 by copybara-service bot Loading…
@danielsuo
[Pallas:MGPU] Add support for clusters and TMA multicast in the Hopper matmul
#31782 opened Sep 12, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Add support for collective copies in pipelining helpers
#31781 opened Sep 12, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Allow changing the number of arrivals on cluster barriers
#31780 opened Sep 12, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Allow varying the arrival count on cluster barriers
#31779 opened Sep 12, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Add tests for the newly added tuning options
#31776 opened Sep 12, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Add Pallas WG lowering rule for the load primitive.
#31775 opened Sep 12, 2025 by copybara-service bot Loading…
[ROCm] Support lowering through PJRT_Triton_Extension
#31768 opened Sep 12, 2025 by amd-jianli12 Loading…
1
[Mosaic:TPU] Push tpu.erase_memref_layout through memref.cast
#31762 opened Sep 12, 2025 by copybara-service bot Loading…
[JAX] Add compiled program call_location logging support to JAX
#31761 opened Sep 11, 2025 by copybara-service bot Loading…
Add remote_bytes_transferred to Pallas CostEstimate dataclass.
#31760 opened Sep 11, 2025 by copybara-service bot Loading…
Reverts e685211bb21b2939590cd353c4949a3226fc391d
#31758 opened Sep 11, 2025 by copybara-service bot Loading…
Re-land https://github.com/jax-ml/jax/pull/31716
#31751 opened Sep 11, 2025 by copybara-service bot Loading…
@jakevdp
[Mosaic GPU][NFC] Always use pre-order traversal.
#31747 opened Sep 11, 2025 by copybara-service bot Loading…
[jaxlib] Do not use the deprecated constructor of absl::*MutexLock
#31740 opened Sep 11, 2025 by copybara-service bot Loading…
@superbobry
Integrate Triton up to
#31739 opened Sep 11, 2025 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.