Skip to content

Pull requests: jax-ml/jax

New pull request New
545 Open 22,543 Closed
545 Open 22,543 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

[xla:ffi] Remove unused UserData ownership forwarding to ffi::ExecutionContext
#33051 opened Oct 31, 2025 by copybara-service bot Loading…
@ezhulenev
[mutable-arrays] make remat work with external refs pull ready Ready for copybara import and testing
#33050 opened Oct 31, 2025 by mattjj Loading…
@mattjj
2
[Pallas] Better error message for grid/dimension semantics mismatch.
#33049 opened Oct 31, 2025 by copybara-service bot Loading…
@justinjfu
Add Pallas core_map guide
#33048 opened Oct 31, 2025 by IvyZX Loading…
30
[JaxMonitoring] Surface unregister listener functions as external APIs.
#33047 opened Oct 31, 2025 by copybara-service bot Loading…
Add tests for transformations of sharded Hijax values
#33045 opened Oct 31, 2025 by samanklesaria • Draft
Support overrideable jax coordinator port and modify default port for Cloud TPU jax.distributed.initialize()
#33044 opened Oct 31, 2025 by copybara-service bot Loading…
Fix handling of pytrees within jnp.array and jnp.asarray.
#33042 opened Oct 31, 2025 by copybara-service bot Loading…
@jakevdp
Relax TpuInfo.is_matmul_supported for TPU7x (Ironwood) to match flexibility already built in HW.
#33040 opened Oct 31, 2025 by copybara-service bot Loading…
[Mosaic] Relax the constraint of doubling contracting size.
#33039 opened Oct 31, 2025 by copybara-service bot Loading…
Delete old Shard Map import code path that isn't needed anymore.
#33037 opened Oct 31, 2025 by copybara-service bot Loading…
[Mosaic GPU] Fix lowering of InitializeBarrierOp.
#33036 opened Oct 31, 2025 by copybara-service bot Loading…
Remove ExecuteShardedArg.
#33034 opened Oct 31, 2025 by copybara-service bot Loading…
@hawkinsp
[MLIR] Fix construction of i1 attributes from numpy arrays and add tests.
#33027 opened Oct 31, 2025 by copybara-service bot Loading…
add hijax coverage tests kokoro:force-run pull ready Ready for copybara import and testing
#33021 opened Oct 31, 2025 by cgarciae Loading…
1
wip remat3 pull ready Ready for copybara import and testing
#33018 opened Oct 30, 2025 by mattjj Loading…
Make k param of jnp.eye non-static.
#33012 opened Oct 30, 2025 by infalmo Loading…
@jakevdp
7
Remove Cholesky forward compatibility path.
#33007 opened Oct 30, 2025 by copybara-service bot Loading…
@hawkinsp
[Pallas:MGPU] Add initial support for TMEM under WG semantic.
#33006 opened Oct 30, 2025 by copybara-service bot Loading…
[mpmd] Add custom pipeline API for user-defined scheduling and merging
#33002 opened Oct 30, 2025 by copybara-service bot Loading…
[Mosaic TPU] No longer implicitly trigger single-buffering when running out of VMEM
#32999 opened Oct 30, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU][NFC] Simplify merging tiling multiples.
#32995 opened Oct 30, 2025 by copybara-service bot Loading…
[Mosaic TPU] Support more general i1 relayout which usually comes from masking. Before this change, we will pick too small sublane tiling for packed types during relayout insertion, e.g., choose (1, 128) for i8.
#32989 opened Oct 30, 2025 by copybara-service bot Loading…
[Mosaic TPU] Add pack/unpack elementwise operations.
#32988 opened Oct 29, 2025 by copybara-service bot Loading…
Test hermetic C++ container
#32985 opened Oct 29, 2025 by copybara-service bot Loading…
@quoctruong
ProTip! Mix and match filters to narrow down what you’re looking for.