Skip to content

Pull requests: jax-ml/jax

New pull request New
664 Open 21,014 Closed
664 Open 21,014 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

[Mosaic GPU][NFC] Always set exact to False in tmem_alloc and tmem_dealloc lowerings.
#31431 opened Aug 29, 2025 by copybara-service bot Loading…
[Mosaic GPU] Remove the exact attribute from tmem_alloc.
#31430 opened Aug 29, 2025 by copybara-service bot Loading…
[Mosaic GPU] Only infer strided layout for vector types.
#31429 opened Aug 29, 2025 by copybara-service bot Loading…
Make copy_to_host_async a no-op and device_get raise a helpful error for non-addressable arrays.
#31427 opened Aug 29, 2025 by copybara-service bot Loading…
@emilyfertig
[Pallas TPU] Add check to ensure that DMA arguments are Refs.
#31424 opened Aug 28, 2025 by copybara-service bot Loading…
@justinjfu
Add common helpers index_dtype_for_shape and index_dtype_for_dim.
#31420 opened Aug 28, 2025 by copybara-service bot Loading…
@hawkinsp
[Mosaic:TPU] Better tpu.memref_reshape support for packed types by improving the reshaped layout inference
#31418 opened Aug 28, 2025 by copybara-service bot Loading…
Add remote_bytes_transferred to Pallas CostEstimate dataclass.
#31415 opened Aug 28, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add lowering rule for tmem_layout_cast.
#31406 opened Aug 28, 2025 by copybara-service bot Loading…
[jax2tf] Remove most of the code for jax2tf with non-native serialization
#31405 opened Aug 28, 2025 by copybara-service bot Loading…
@gnecula
[jax2tf] Remove most of the code for jax2tf with non-native serialization pull ready Ready for copybara import and testing
#31404 opened Aug 28, 2025 by gnecula Loading…
@gnecula
Package Sdy MPMD into Jax MLIR wheel
#31402 opened Aug 28, 2025 by copybara-service bot Loading…
[AArch64 CPU] Cast BFloat16 Gelu to F32
#31401 opened Aug 28, 2025 by Ryo-not-rio Loading…
Reverts 402a78661747fc1fde0bc37200996b26eea1c9da
#31394 opened Aug 27, 2025 by copybara-service bot Loading…
Run TPU 7x int->float matmul canonicalization before the generic mixed-dtype canonicalization.
#31392 opened Aug 27, 2025 by copybara-service bot Loading…
Remove old ROCm build code
#31381 opened Aug 27, 2025 by charleshofer Loading…
Add uv Installation Instructions to Documentation (GH#29451)
#31377 opened Aug 27, 2025 by Aniketsy Loading…
Clarify input_output_aliases indexing in pallas_call.
#31376 opened Aug 27, 2025 by copybara-service bot Loading…
[gpu] Always use the Jacobi algorithm for lax.linalg.eigh
#31375 opened Aug 27, 2025 by jaro-sevcik Loading…
1
2
Iterate in reverse order for TPU7x to preferentially map the innermost physical axis for higher bandwidth.
#31373 opened Aug 27, 2025 by copybara-service bot Loading…
[Mosaic] Support any N-D permutations.
#31372 opened Aug 27, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add layout inference rule for AsyncLoadTmemOp and AsyncStoreTmemOp.
#31371 opened Aug 27, 2025 by copybara-service bot Loading…
[Export] Add AllocateBuffer as a supported custom call.
#31369 opened Aug 27, 2025 by copybara-service bot Loading…
@bchetioui
[NFC] Refactor the collective matmul
#31367 opened Aug 27, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic] Add support for S32 cross-lane reduction.
#31359 opened Aug 27, 2025 by copybara-service bot Loading…
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.