Skip to content

Pull requests: jax-ml/jax

New pull request New
553 Open 21,768 Closed
553 Open 21,768 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

Replace IFRT futures/promises with TSL futures/promises
#32162 opened Sep 28, 2025 by copybara-service bot Loading…
Use triple overloads of TargetRegistry functions
#32161 opened Sep 28, 2025 by boomanaiden154 Loading…
Expose jax.export at top-level for consistency with other submodules
#32160 opened Sep 28, 2025 by athulya-anil Loading…
1
buffer_callback doesn't require devices during lowering. So remove the registration. Fixes https://github.com/jax-ml/jax/issues/32149
#32158 opened Sep 27, 2025 by copybara-service bot Loading…
1
@yashk2810
Update run-bazel-test-cuda-py-import job to run multiaccelerator tests on L4 GPUs.
#32155 opened Sep 27, 2025 by copybara-service bot Loading…
[Pallas][Mosaic GPU] Add docstrings for plgpu.kernel.
#32153 opened Sep 27, 2025 by copybara-service bot Loading…
@justinjfu
[hijax] vmap is not a primitive, it is sophisticated, hi-class pull ready Ready for copybara import and testing
#32152 opened Sep 26, 2025 by mattjj Loading…
@mattjj
add aval_property.setter
#32151 opened Sep 26, 2025 by cgarciae Loading…
[pmap] Fix boolean logic typo in shard_map implementation.
#32146 opened Sep 26, 2025 by copybara-service bot Loading…
@danielsuo
[haiku] Update for jax_pmap_shmap_merge=True.
#32144 opened Sep 26, 2025 by copybara-service bot Loading…
@danielsuo
[Pallas] Add default index map instead of keeping it None
#32143 opened Sep 26, 2025 by copybara-service bot Loading…
@sharadmv
[mutable arrays] remove internal aliases pull ready Ready for copybara import and testing
#32142 opened Sep 26, 2025 by jakevdp Loading…
@jakevdp
Remove notap flag from B200 pallas tests
#32141 opened Sep 26, 2025 by copybara-service bot Loading…
Use common instead of build in .bazelrc file.
#32137 opened Sep 26, 2025 by copybara-service bot Loading…
[Mosaic GPU] Move _is_mma_layout to inference_utils and reuse it.
#32134 opened Sep 26, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Improve the benchmark harness for Blackwell matmul
#32133 opened Sep 26, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Add a Blackwell matmul tutorial documentation
#32132 opened Sep 26, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add a Divides constraint type.
#32131 opened Sep 26, 2025 by copybara-service bot Loading…
[Pallas/Mosaic GPU] Expand the lowering capabilities of broadcast_in_dim with warpgroup semantics.
#32129 opened Sep 26, 2025 by copybara-service bot Loading…
@bchetioui
[Pallas/Mosaic GPU] Conditionally construct a fa_m64_collective_layout when inferring layouts for broadcast_in_dim.
#32128 opened Sep 26, 2025 by copybara-service bot Loading…
@bchetioui
[Pallas][Mosaic GPU] Add work stealing to blackwell matmul kernel.
#32126 opened Sep 26, 2025 by copybara-service bot Loading…
@justinjfu
[JAX] Add compiled program call_location logging support to JAX
#32124 opened Sep 25, 2025 by copybara-service bot Loading…
[Mosaic] Expand arith.select canonicalization
#32122 opened Sep 25, 2025 by copybara-service bot Loading…
[CI] Ignore: Tune the B200 tests per accelerator CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#32121 opened Sep 25, 2025 by MichaelHudgins • Draft
[Mosaic TPU] Support scalar input-output-alias. Relax the condition in ShapeVerifier::HandleCustomCall to allow scalar operand live on S(6) while the result has no memory space assigned.
#32116 opened Sep 25, 2025 by copybara-service bot Loading…
ProTip! no:milestone will show everything without a milestone.