Skip to content

Pull requests: jax-ml/jax

New pull request New
552 Open 25,110 Closed
552 Open 25,110 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

[Mosaic GPU] Lift restriction on tcgen05.mma sparse metadata WG lowering.
#35957 opened Mar 17, 2026 by copybara-service bot Loading…
[Mosaic GPU] Add Pallas WG lowering for async_copy_sparse_metadata_to_tmem.
#35956 opened Mar 17, 2026 by copybara-service bot Loading…
[Mosaic GPU] Add TMEM aliasing support for WG semantics.
#35955 opened Mar 17, 2026 by copybara-service bot Loading…
Bump the maximum CUDA PJRT plugin wheel size.
#35954 opened Mar 17, 2026 by copybara-service bot Loading…
Add temporary backward compatbility for earlier method of specialization.
#35952 opened Mar 17, 2026 by copybara-service bot Loading…
Suppressed pytype errors in uses of MLIR APIs
#35951 opened Mar 17, 2026 by copybara-service bot Loading…
@superbobry
[Shardy MPMD] Add support for unreduced axes in MPMD sharding.
#35949 opened Mar 17, 2026 by copybara-service bot Loading…
@emilyfertig
[pallas] rewrite pull_block_spec rule to work in terms of block_index transforms
#35948 opened Mar 17, 2026 by copybara-service bot Loading…
@levskaya
unreduced/reduced in ShapedArray, NamedSharding can only talk about Explicit axes.
#35945 opened Mar 16, 2026 by copybara-service bot Loading…
@yashk2810
Add ROCm xdist device pinning in conftest.py
#35944 opened Mar 16, 2026 by gulsumgudukbay Loading…
@hawkinsp
5
internal change
#35940 opened Mar 16, 2026 by copybara-service bot Loading…
@tomhennigan
jnp.arange: avoid constant-folding when using step
#35933 opened Mar 16, 2026 by copybara-service bot Loading…
@jakevdp
Rocm/s3 wheel downloads AMD GPU Issues pertaining to AMD GPUs (ROCM) CI kokoro:force-run pull ready Ready for copybara import and testing
#35932 opened Mar 16, 2026 by mminutoli Loading…
3
[Mosaic] Enable internal fptoi/itofp ops for scalar types
#35931 opened Mar 16, 2026 by copybara-service bot Loading…
[Mosaic GPU] Add WarpMapOp to the dialect.
#35926 opened Mar 16, 2026 by copybara-service bot Loading…
Produce and upload junit XML files in GitHub Actions test workflows as GitHub Actions artifacts.
#35923 opened Mar 16, 2026 by copybara-service bot Loading…
@danielsuo
Introduce a new hypothesis_test_util module.
#35922 opened Mar 16, 2026 by copybara-service bot Loading…
@hawkinsp
[ROCm] Skip test_tridiagonal_solve_grad0 due to rocSparse numerical issue
#35917 opened Mar 16, 2026 by AratiGanesh Loading…
1
Add memory-efficient fused linear cross-entropy loss (solves issue #35906)
#35915 opened Mar 15, 2026 by Coffeempty Loading…
1
10
Add unroll parameter to lax.while_loop.
#35913 opened Mar 14, 2026 by carlosgmartin Loading…
1
3
[pallas-triton] Relax compute capability requirement in pallas-triton tests.
#35910 opened Mar 14, 2026 by copybara-service bot Loading…
@rdyro
Add lax.associative_reduce.
#35904 opened Mar 14, 2026 by carlosgmartin Loading…
1
5
Identify backend type of kernels from MLIR,
#35903 opened Mar 14, 2026 by copybara-service bot Loading…
[ROCm] Enable hipBLASLt in ROCm pytest CI
#35902 opened Mar 14, 2026 by magaonka-amd Loading…
3
[Mosaic] Add tpu.sfence
#35901 opened Mar 13, 2026 by copybara-service bot Loading…
ProTip! Updated in the last three days: updated:>2026-03-14.