Skip to content

Pull requests: jax-ml/jax

New pull request New
657 Open 23,345 Closed
657 Open 23,345 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

feat: Add weights parameter to quantile and percentile functions (#32647)
#34084 opened Dec 22, 2025 by C-you-know Loading…
1
6
[NFC] Remove obsolete version check in absl_cpp_logging_test.
#34081 opened Dec 22, 2025 by copybara-service bot Loading…
docs: clarify JAX intro video link in beginner guide
#34078 opened Dec 21, 2025 by NAS-24 Loading…
4
Fix/checkpoint static argnames
#34077 opened Dec 21, 2025 by Indrapal-70 Loading…
6
Add file identifier and extension to export schema
#34075 opened Dec 20, 2025 by lutzroeder Loading…
4
Add wheel sources test.
#34071 opened Dec 19, 2025 by copybara-service bot Loading…
@ybaturina
1
Remove some more stale version guards.
#34069 opened Dec 19, 2025 by copybara-service bot Loading…
@hawkinsp
DO_NOT_SUBMIT: Testing TAP
#34068 opened Dec 19, 2025 by copybara-service bot Loading…
@mwhittaker
[dep] finalize a number of deprecations for JAX v0.9.0
#34067 opened Dec 19, 2025 by copybara-service bot Loading…
@jakevdp
Add support for specifying NCCL version with nccl_version flag
#34063 opened Dec 19, 2025 by ezhulenev Loading…
4
Fix the licenses in jax/BUILD file as a follow-up after [this change](https://github.com/jax-ml/jax/commit/7bc428b8a0204b384e4236181ebe5bf1cac02b5c).
#34057 opened Dec 19, 2025 by copybara-service bot Loading…
@ybaturina
1
Internal-only change
#34056 opened Dec 19, 2025 by copybara-service bot Loading…
@belitskiy
[pmap] Update migrate pmap docs with additional user errors and performance notes.
#34053 opened Dec 19, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Add support for TMA gather/scatter in WG semantics.
#34052 opened Dec 19, 2025 by copybara-service bot Loading…
Add programmatic profiling test with session_id support in JAX.
#34051 opened Dec 19, 2025 by copybara-service bot Loading…
Update convertReshardToShardedToUnreduced to generate replicated-to-unreduced + sharded-to-unreduced.
#34038 opened Dec 19, 2025 by copybara-service bot Loading…
Update the lock files post the 0.8.2 release.
#34035 opened Dec 18, 2025 by copybara-service bot Loading…
@belitskiy
[Autotuner] Combine all* autotuners in one backend-agnostic pass.
#34034 opened Dec 18, 2025 by copybara-service bot Loading…
Add the targets needed for Bazel tests but not included in JAX wheels in wheel_additives.
#34031 opened Dec 18, 2025 by copybara-service bot Loading…
@ybaturina
1
[Pallas/Mosaic GPU] Make Mosaic GPU the default lowering backend for Pallas on GPU.
#34027 opened Dec 18, 2025 by copybara-service bot Loading…
@bchetioui
[Docs] lax.conv_shape_tuple: Docstring and clarified shape assumptions and padding
#34026 opened Dec 18, 2025 by Prakharprasun Loading…
5
Automated Code Change
#34019 opened Dec 18, 2025 by copybara-service bot Loading…
[Pallas] Add trivial batching rule for with_memory_space_constraint
#34016 opened Dec 18, 2025 by copybara-service bot Loading…
@sharadmv
Test nvshmem changes.
#34011 opened Dec 17, 2025 by copybara-service bot Loading…
@ybaturina
1
Update XLA dependency to use revision http://github.com/openxla/xla/commit/f335d3d6cfc367dc32de84a76c62f2fc9a08d22e
#34008 opened Dec 17, 2025 by copybara-service bot Loading…
@belitskiy
ProTip! Filter pull requests by the default branch with base:main.