Skip to content

Pull requests: jax-ml/jax

New pull request New
619 Open 23,121 Closed
619 Open 23,121 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

Automated Code Change
#33789 opened Dec 7, 2025 by copybara-service bot Loading…
Add LoadedExecutable.serialize() to match Executable.serialize().
#33786 opened Dec 6, 2025 by copybara-service bot Loading…
@pschuh
improve tracer reprs to show a lil concrete info kokoro:force-run pull ready Ready for copybara import and testing
#33785 opened Dec 6, 2025 by mattjj Loading…
3
Allow overriding the XLA commit to use in Continuous workflow
#33784 opened Dec 6, 2025 by copybara-service bot Loading…
@nitins17
Export pyproject.toml in the BUILD file
#33783 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
[jax.collect_profile] Allow arbitrary options to be passed to XProf
#33782 opened Dec 5, 2025 by copybara-service bot Loading…
[Pallas TPU] Allow closed over scalars in core_map code
#33781 opened Dec 5, 2025 by copybara-service bot Loading…
@sharadmv
Add collective to async_copy_scales_to_tmem.
#33780 opened Dec 5, 2025 by copybara-service bot Loading…
[Pallas SC] Allow semaphores to be returned by SCS kernels
#33779 opened Dec 5, 2025 by copybara-service bot Loading…
@sharadmv
[pallas:mosaic] pltpu.emit_pipeline now accepts block specs in HBM
#33776 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
Change PjRt to use new copy of coordination service.
#33775 opened Dec 5, 2025 by copybara-service bot Loading…
@mwhittaker
Move rules_ml_toolchain up due to request move default rules_ml_toolchain initialization to @xla//:workspace4.bzl
#33774 opened Dec 5, 2025 by copybara-service bot Loading…
1
Colocated python perf optimization
#33772 opened Dec 5, 2025 by copybara-service bot Loading…
Add initial blackwell mxfp4 blockscale matmul implementation.
#33770 opened Dec 5, 2025 by copybara-service bot Loading…
[pallas:mosaic] Access memory spaces via pltpu directly
#33768 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
[pallas:sc] pltpu.emit_pipeline now supports 1D-tiled block specs
#33767 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
[mosaic] Added documentation and a few useful methods to tpu.tiled attribute
#33763 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
Adjust tile_n assert to accommodate 2 cta in tcgen05 blockscale mma.
#33760 opened Dec 5, 2025 by copybara-service bot Loading…
[XLA:MGPU] Port Tiling to C++.
#33759 opened Dec 5, 2025 by copybara-service bot Loading…
[export] Add backwards compatibility tests for serialized exports pull ready Ready for copybara import and testing
#33758 opened Dec 5, 2025 by gnecula Loading…
@gnecula
3
register rocm to block_scaled_dot lowering path
#33757 opened Dec 5, 2025 by ScXfjiang Loading…
6
[Mosaic GPU] Add support for a limited class of reshapes for tiled layouts
#33756 opened Dec 5, 2025 by copybara-service bot Loading…
@apaszke
WIP: Indexing refactor
#33752 opened Dec 5, 2025 by jakevdp • Draft
@jakevdp
3
Merge repr and str of array in one function instead of duplicating conditions
#33750 opened Dec 5, 2025 by copybara-service bot Loading…
@yashk2810
ref dup and merge operations, for parallel writers
#33749 opened Dec 5, 2025 by mattjj • Draft
6
ProTip! Add no:assignee to see everything that’s not assigned.