Skip to content

Pull requests: jax-ml/jax

New pull request New
688 Open 26,940 Closed
688 Open 26,940 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

[pallas:mgpu] Expose TMA_INDICES_4 layout to gather/scatter <16 rows.
#38175 opened Jun 3, 2026 by copybara-service Bot Loading…
Fix spurious KeyReuseError under vmap over a zero-sized key array
#38174 opened Jun 3, 2026 by Lawson-Darrow Loading…
1
Move dtypes.canonicalize_value to C++.
#38171 opened Jun 3, 2026 by copybara-service Bot Loading…
@hawkinsp
ci: Add OneAPI artifact build script and auditwheel support
#38170 opened Jun 3, 2026 by ashahba Contributor Loading…
3 tasks done
6
put ct avals on GradAccums (and UndefinedPrimals), not primal avals kokoro:force-run pull ready Ready for copybara import and testing
#38168 opened Jun 3, 2026 by mattjj Collaborator Loading…
@mattjj
7
diffbase needs this section.
#38167 opened Jun 3, 2026 by copybara-service Bot Loading…
Switch api.py AD functions from linear util to FlatTrees pull ready Ready for copybara import and testing
#38166 opened Jun 3, 2026 by dougalm Collaborator Loading…
@dougalm
4
Support prefetched inputs in Pallas pipelines
#38165 opened Jun 3, 2026 by copybara-service Bot Loading…
Test core split for TPU7x.
#38162 opened Jun 3, 2026 by copybara-service Bot Loading…
@belitskiy
Make JAX axis names deterministic across executions.
#38161 opened Jun 3, 2026 by copybara-service Bot Loading…
Add Thor SM arch to CUDA build
#38155 opened Jun 3, 2026 by aybchan Member Loading…
2
[pallas] Fix Pallas Mosaic 64-bit mode loop indexing
#38154 opened Jun 3, 2026 by copybara-service Bot Loading…
@gnecula
[Mosaic:GPU] Remove NVSHMEM usage from Mosaic XLA custom kernel and tests.
#38153 opened Jun 3, 2026 by copybara-service Bot Loading…
[Mosaic TPU] Support static lane shift of pltpu.roll with stride > 0 for packed types, where the lane dimension should be aligned to 128.
#38147 opened Jun 3, 2026 by copybara-service Bot Loading…
[ROCm] Enable HLO module transform registration for GPU backends AMD GPU Issues pertaining to AMD GPUs (ROCM)
#38142 opened Jun 2, 2026 by mminutoli Contributor Loading…
2
[JAX] Remove .count for jax._src.core.Var.
#38138 opened Jun 2, 2026 by copybara-service Bot Loading…
@hawkinsp
[Mosaic GPU] Use MOSAIC_GPU_CUDA_ROOT to define the path to CUDA tools.
#38136 opened Jun 2, 2026 by copybara-service Bot Loading…
@bchetioui
[Mosaic GPU] WGxWarp semantics for async_copy_sparse_metadata_to_tmem.
#38135 opened Jun 2, 2026 by copybara-service Bot Loading…
[OneAPI GPU] Add OneAPI (SYCL) GPU Runtime Functions.
#38134 opened Jun 2, 2026 by kranipa Contributor Loading…
5
[Mosaic GPU] Add a property-based test for TMEM column slicing.
#38131 opened Jun 2, 2026 by copybara-service Bot Loading…
[Mosaic GPU] Enable column slicing for TMEM layouts with varying row dimensions.
#38129 opened Jun 2, 2026 by copybara-service Bot Loading…
[pallas:sc] Added support for tpu.scan to vector layout passes
#38126 opened Jun 2, 2026 by copybara-service Bot Loading…
@superbobry
Fix jnp.vectorize crash when None is at index 0
#38117 opened Jun 2, 2026 by Kymi808 Loading…
@jakevdp
6
Plumb tiling info from jax through stable hlo for xla's LayoutConstraint custom-call.
#38109 opened Jun 1, 2026 by copybara-service Bot Loading…
Add addlicense pre-commit hook to check for license headers.
#38104 opened Jun 1, 2026 by copybara-service Bot Loading…
ProTip! Follow long discussions with comments:>50.