Skip to content

Pull requests: jax-ml/jax

New pull request New
699 Open 26,482 Closed
699 Open 26,482 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

[export] Handle AMD device kinds in get_platform_from_device_kind kokoro:force-run pull ready Ready for copybara import and testing
#37634 opened May 13, 2026 by zahiqbal Contributor Loading…
6
[CI] Allow H100+B200 workflow dispatch to specify XLA commit, test targets, jobs.
#37632 opened May 13, 2026 by copybara-service Bot Loading…
@belitskiy
Introduce layout_constraint mosaic op.
#37629 opened May 13, 2026 by copybara-service Bot Loading…
@yashk2810
[Mosaic GPU][NFC] Enforce exhaustive Constant checks in constraints.py.
#37626 opened May 13, 2026 by copybara-service Bot Loading…
[Mosaic:GPU] Reland: Use NCCL API on the host to collective peer addresses.
#37623 opened May 13, 2026 by copybara-service Bot Loading…
Upgrade bazel version to 8.7.0
#37622 opened May 13, 2026 by copybara-service Bot Loading…
@meteorcloudy
[Pallas:TPU] Use the new late loop unroller pass in Mosaic when available
#37617 opened May 13, 2026 by copybara-service Bot Loading…
@apaszke
Enable Bzlmod for Jax
#37615 opened May 13, 2026 by copybara-service Bot Loading…
@meteorcloudy
Make stack and unstack jax primitives.
#37600 opened May 13, 2026 by copybara-service Bot Loading…
Add jax.scipy.special.comb
#37599 opened May 13, 2026 by bej9 Contributor Loading…
@jakevdp
8
Add debug_leaked_clients_on_clear_backends which allows checking
#37598 opened May 13, 2026 by copybara-service Bot Loading…
@pschuh
Add pl.concat_ref as a transformed ref for Pallas DMA calls.
#37596 opened May 12, 2026 by copybara-service Bot Loading…
@IvyZX
Add poison_buffers to emit_pipeline to poison intermediate buffers for testing.
#37593 opened May 12, 2026 by copybara-service Bot Loading…
@sharadmv
Fix various hijax lowering issues
#37592 opened May 12, 2026 by copybara-service Bot Loading…
@cgarciae
[Pallas] Expose additional classes from pallas.core package into jax.extend.pallas
#37584 opened May 12, 2026 by copybara-service Bot Loading…
[pallas] Added a batching rule for mpmd_map_p
#37574 opened May 12, 2026 by copybara-service Bot Loading…
@superbobry
Add push_block_spec support for lax.concatenate_p in Pallas fuser.
#37569 opened May 12, 2026 by copybara-service Bot Loading…
[SDY] Make the sdy_version argument mandatory in StableHlo serialization APIs
#37568 opened May 12, 2026 by copybara-service Bot Loading…
Add jax.scipy.special.boxcox and boxcox1p
#37564 opened May 11, 2026 by Saumay Contributor Loading…
@jakevdp
5
Exclude --xla_tpu_compiler_variant from JAX cache keys
#37561 opened May 11, 2026 by copybara-service Bot Loading…
Playing with imports.
#37560 opened May 11, 2026 by copybara-service Bot Loading…
@mwhittaker
Include resolved XLA TPU compiler version in JAX cache key
#37558 opened May 11, 2026 by copybara-service Bot Loading…
Bump the pip group across 1 directory with 2 updates dependencies Pull requests that update a dependency file python Pull requests that update Python code
#37553 opened May 11, 2026 by dependabot Bot Loading…
[Mosaic GPU][NFC] Introduce AlwaysTrue constraint.
#37551 opened May 11, 2026 by copybara-service Bot Loading…
Use native Python slices instead of pl.ds where possible
#37550 opened May 11, 2026 by copybara-service Bot Loading…
@superbobry
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.