Skip to content

Pull requests: jax-ml/jax

New pull request New
566 Open 22,079 Closed
566 Open 22,079 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

fix(docs): regularized incomplete beta function (jax.scipy.special.betainc) formula corrected in docstring
#32539 opened Oct 11, 2025 by Qazalbash Loading…
Automated Code Change
#32538 opened Oct 11, 2025 by copybara-service bot Loading…
Bump rules_ml_toolchain version to integrate @cuda_nvrtc header.
#32534 opened Oct 10, 2025 by copybara-service bot Loading…
@ybaturina
1
Deprecate jax.experimental.shard_map kokoro:force-run pull ready Ready for copybara import and testing
#32533 opened Oct 10, 2025 by jakevdp Loading…
@jakevdp
Remove use of vector.splat
#32530 opened Oct 10, 2025 by copybara-service bot Loading…
[Pallas/interpreter] Move SharedMemory out of the main interpreter source file.
#32529 opened Oct 10, 2025 by copybara-service bot Loading…
[Mosaic GPU][NFC] Use dataclasses.fields to gather operands for CustomPrimitiveOp.
#32523 opened Oct 10, 2025 by copybara-service bot Loading…
Add missing function handler to nvml impl lib and use it
#32521 opened Oct 10, 2025 by copybara-service bot Loading…
[Mosaic] Relax the restrictions of small second minor tiling reshape.
#32518 opened Oct 10, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add a basic test for running collectives on Torch tensors kokoro:force-run pull ready Ready for copybara import and testing
#32514 opened Oct 10, 2025 by apaszke Loading…
4
[pallas:sc] Fixed the lowering of a dynamically sized gather where the indices ref has more than one index
#32513 opened Oct 10, 2025 by copybara-service bot Loading…
@superbobry
[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…
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
3
Update rules_ml_toolchain and set LLVM version to 21.
#32489 opened Oct 9, 2025 by copybara-service bot Loading…
1
[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
ProTip! Add no:assignee to see everything that’s not assigned.