Skip to content

Pull requests: jax-ml/jax

New pull request New
565 Open 22,052 Closed
565 Open 22,052 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] Implement a basic reduce-scatter collective using multimem.ld_reduce
#32511 opened Oct 10, 2025 by copybara-service bot Loading…
@apaszke
Document that negative low/high padding is allowed for lax.pad.
#32509 opened Oct 10, 2025 by copybara-service bot Loading…
@hawkinsp
[Mosaic GPU] Support profiling for WG semantic.
#32507 opened Oct 10, 2025 by copybara-service bot Loading…
[pallas:sc] Avoid putting tracers into params of dma_* primitives
#32506 opened Oct 10, 2025 by copybara-service bot Loading…
@superbobry
Add "relaxed" mode for "jax_numpy_dtype_promotion".
#32504 opened Oct 10, 2025 by copybara-service bot Loading…
@chr1sj0nes
Check that exceptions are thrown when invalid types are passed to promote_types.
#32503 opened Oct 10, 2025 by copybara-service bot Loading…
@chr1sj0nes
reshape->store into strided_store rewrite
#32502 opened Oct 10, 2025 by copybara-service bot Loading…
[Mosaic] Increase the number of partial accumulators used in multi_reduction to expose more ILP. When in shape invariant mode, fall back to use sequential reduction with single partial accumulator.
#32500 opened Oct 10, 2025 by copybara-service bot Loading…
[mutable-arrays] fix batching bugs in ref primitives kokoro:force-run pull ready Ready for copybara import and testing
#32499 opened Oct 10, 2025 by mattjj Loading…
1
@mattjj
1
fix pjit dce effect indexing, cf #32434
#32496 opened Oct 10, 2025 by mattjj Loading…
1
prevent reduce jvp from reordering reduction
#32495 opened Oct 10, 2025 by mattjj Loading…
1
2
[hijax] add a QArray example using new-style hijax primitives kokoro:force-run pull ready Ready for copybara import and testing
#32493 opened Oct 9, 2025 by mattjj Loading…
Switch JAX to use an FFI call for cholesky decompositions rather than the Cholesky HLO on GPU.
#32492 opened Oct 9, 2025 by copybara-service bot Loading…
@hawkinsp
Switch the default nonsymmetric eigendecomposition implementation to use cusolver.
#32490 opened Oct 9, 2025 by copybara-service bot Loading…
1
@hawkinsp
Update rules_ml_toolchain and set LLVM version to 21.
#32489 opened Oct 9, 2025 by copybara-service bot Loading…
1
[Pallas:MGPU] Add a collective matmul tutorial documentation
#32488 opened Oct 9, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic] Add sign_extended attribute to tpu::UnpackSubelementsOp.
#32477 opened Oct 9, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Add ReducedLayout support to inline_mgpu params
#32471 opened Oct 9, 2025 by copybara-service bot Loading…
DNM [vjp3] Make value_and_grad always use f_vjp.with_refs(*accums)
#32470 opened Oct 9, 2025 by gspschmid • Draft
Add MPMD to IFRT IR lowerings to jaxlib._sdy_mpmd
#32466 opened Oct 9, 2025 by copybara-service bot Loading…
[Mosaic GPU] Support profiling for WG semantic.
#32465 opened Oct 9, 2025 by copybara-service bot Loading…
Fix dtype promotion issue in scalar indexing pull ready Ready for copybara import and testing
#32460 opened Oct 8, 2025 by jakevdp Loading…
1
@jakevdp
[JAX] Add jax_check_ifrt_user_context config and enable it for JAX tests
#32458 opened Oct 8, 2025 by copybara-service bot Loading…
Use a single fixed channel ID for all collectives emitted by JAX.
#32456 opened Oct 8, 2025 by copybara-service bot Loading…
@hawkinsp
Automated Code Change
#32444 opened Oct 8, 2025 by copybara-service bot Loading…
ProTip! Adding no:label will show everything without a label.