Skip to content

Pull requests: jax-ml/jax

New pull request New
564 Open 22,790 Closed
564 Open 22,790 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

Transposed ragged dot kernel on Blackwell GPU
#33352 opened Nov 16, 2025 by copybara-service bot Loading…
fix ref effect indexing bugs pull ready Ready for copybara import and testing
#33351 opened Nov 16, 2025 by mattjj • Draft
@mattjj
[Mosaic TPU] Add a new canonicalization pattern for tpu::ReshapeOp, targeting at folding last two dims. This could optimize existing reshapes and also support more cases, e.g., when the src minor dim is not divisible by lane count.
#33346 opened Nov 15, 2025 by copybara-service bot Loading…
[Pallas/Fuser] Add basic for some simple output concatenation fusion
#33345 opened Nov 15, 2025 by copybara-service bot Loading…
@sharadmv
[Mosaic] Internal change.
#33343 opened Nov 14, 2025 by copybara-service bot Loading…
Test new rules_ml_toolchain version for CUDNN10.
#33339 opened Nov 14, 2025 by copybara-service bot Loading…
@ybaturina
1
[Pallas:MGPU] Fix plgpu.inline_mgpu support for nested blocks.
#33337 opened Nov 14, 2025 by copybara-service bot Loading…
[Pallas:SC] Enable load_gather_masked tests with TC tiling
#33336 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:SC] Test with both SC and TC tiling
#33335 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[mpmd] Modify tests to also test custom pipeline API path
#33334 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Add custom pipeline API for user-defined scheduling and merging
#33333 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Move pipeline scheduling types and conversion functions to OSS
#33332 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Fix nanobind compatibility and add OPTIMIZE enum value
#33331 opened Nov 14, 2025 by copybara-service bot Loading…
Reduce the maximum test parallelism on Windows machines to avoid OOMs
#33330 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[DO NOT SUBMIT] Inject failure into CI
#33327 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Support 1D input and keepdims=False for argmax/argmin.
#33321 opened Nov 14, 2025 by copybara-service bot Loading…
[xla:cpu] Don't hoist small loops with runtime calls
#33319 opened Nov 14, 2025 by copybara-service bot Loading…
@ezhulenev
Add test of documentation coverage pull ready Ready for copybara import and testing
#33316 opened Nov 13, 2025 by jakevdp Loading…
@jakevdp
5
Set xla_tpu_enable_mlir_verifier in tests
#33311 opened Nov 13, 2025 by copybara-service bot Loading…
[Mosaic GPU] Use a recursive tree implementation for vector_concat
#33305 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Tune the packing length for i4 -> f8e4m3fn conversion
#33304 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[mosaic] Made GetMosaicHardwareGeneration public
#33301 opened Nov 13, 2025 by copybara-service bot Loading…
@superbobry
Automated Code Change
#33293 opened Nov 13, 2025 by copybara-service bot Loading…
[pmap] Implement device_put_sharded using NamedSharding.
#33292 opened Nov 13, 2025 by copybara-service bot Loading…
@danielsuo
Enable RBE testing of JAX CUDA 13 wheels by loading CUDA UMD libs from Python code.
#33291 opened Nov 13, 2025 by copybara-service bot Loading…
@ybaturina
1
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.