Skip to content

Pull requests: jax-ml/jax

New pull request New
600 Open 22,999 Closed
600 Open 22,999 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 pvary with pcast in shard_map doc pull ready Ready for copybara import and testing
#33629 opened Dec 1, 2025 by yashk2810 Loading…
1
[Pallas] Fix cost estimate einsum bug.
#33628 opened Dec 1, 2025 by copybara-service bot Loading…
@justinjfu
Update rules_ml_toolchain to disable --@local_config_cuda//cuda:include_cuda_libs by default.
#33627 opened Dec 1, 2025 by copybara-service bot Loading…
@ybaturina
1
Add CI to run ty on jax/_src/dtypes.py
#33626 opened Dec 1, 2025 by jakevdp Loading…
@jakevdp
4
Deprecate jax.lax.pvary in favor of jax.lax.pcast
#33625 opened Dec 1, 2025 by copybara-service bot Loading…
@yashk2810
Remove an apparently stale test skip.
#33624 opened Dec 1, 2025 by copybara-service bot Loading…
@hawkinsp
Replace jax.lax.pvary with jax.lax.pcast(..., to='varying'). No function changes intended. This is just a rename.
#33621 opened Dec 1, 2025 by copybara-service bot Loading…
@yashk2810
Introduce jax.lax.pcast as an API to replace pvary, preduced, reduced_vary_cast and vary_unreduced_cast APIs.
#33619 opened Dec 1, 2025 by copybara-service bot Loading…
@yashk2810
Build JAX artifacts only in CPU pools for RBE configurations.
#33618 opened Dec 1, 2025 by copybara-service bot Loading…
@ybaturina
1
[mosaic] Added a doc comment and a check to tpu::ComputeTileStrides
#33615 opened Dec 1, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Impelment basic support for memref.expand_shape in WG semantics.
#33612 opened Dec 1, 2025 by copybara-service bot Loading…
If LAPACK kernels are linked statically, don't override them from scipy.
#33611 opened Dec 1, 2025 by copybara-service bot Loading…
@hawkinsp
[Pallas:SC] Add support for neg_p, abs_p.
#33610 opened Dec 1, 2025 by copybara-service bot Loading…
Fix links and improve clarity in PRNG documentation
#33608 opened Dec 1, 2025 by partev Loading…
@jakevdp
7
[doc] Correct link to JAX core library
#33606 opened Nov 30, 2025 by partev Loading…
2
Improve numerical stability of BiCGStab on GPU
#33603 opened Nov 30, 2025 by yashwantbezawada Loading…
1
@jakevdp
12
Fix constant folding bug in jnp.arange for non-zero start values
#33602 opened Nov 29, 2025 by yashwantbezawada Loading…
1
@jakevdp
8
Add support for NumPy 2.4 __numpy_dtype__ protocol
#33601 opened Nov 29, 2025 by jakevdp Loading…
@jakevdp
3
Automated Code Change
#33600 opened Nov 29, 2025 by copybara-service bot Loading…
[Mosaic GPU] Implement GPU module unloading for Mosaic GPU custom calls.
#33599 opened Nov 28, 2025 by copybara-service bot Loading…
[pallas:sc] Skip a few tests failing when the compiler uses tiled memrefs
#33598 opened Nov 28, 2025 by copybara-service bot Loading…
@superbobry
[export] Add support for explicit sharding. pull ready Ready for copybara import and testing
#33597 opened Nov 28, 2025 by gnecula Loading…
@gnecula
5
Add programmatic profiling test with session_id support in JAX.
#33596 opened Nov 28, 2025 by copybara-service bot Loading…
Add ppf (inverse CDF) for Gamma and Beta distributions (#20358)
#33589 opened Nov 28, 2025 by VARUN3WARE Loading…
1
8
[Mosaic GPU] Add support for arbitrary bitwidths in 32x32b TMEM transfers
#33580 opened Nov 27, 2025 by copybara-service bot Loading…
@apaszke
ProTip! Add no:assignee to see everything that’s not assigned.