Skip to content

Pull requests: jax-ml/jax

New pull request New
582 Open 19,962 Closed
582 Open 19,962 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:MGPU][Experimental] HLO -> Pallas.
#30122 opened Jul 10, 2025 by copybara-service bot Loading…
[Pallas] Make platforms param configurable for pl.lower_as_mlir.
#30121 opened Jul 10, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Disallow direct reads from TMEM refs
#30119 opened Jul 10, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Expose block-scaled MMA to Pallas on Blackwell
#30116 opened Jul 10, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Align TMEM allocations to 16 bytes
#30115 opened Jul 10, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU][NFC] Add utils to create and check the SMEM memory space attribute.
#30114 opened Jul 10, 2025 by copybara-service bot Loading…
[Mosaic GPU][NFC] Enable the gmem argument to async_copy to be an ir.BlockArgument
#30112 opened Jul 10, 2025 by copybara-service bot Loading…
[Pallas TPU] Refactor casting logic to use bits instead of bytes and allow uint upcasts.
#30103 opened Jul 9, 2025 by copybara-service bot Loading…
@justinjfu
Pallas/Mosaic support host input.
#30102 opened Jul 9, 2025 by copybara-service bot Loading…
Don't export Shardy in MPMD before going to IFRT
#30093 opened Jul 9, 2025 by copybara-service bot Loading…
#sdy Remove MHLO shardings from round-trip export
#30091 opened Jul 9, 2025 by copybara-service bot Loading…
Update Protobuf to 6.31.1
#30089 opened Jul 9, 2025 by copybara-service bot Loading…
remove pjit_p, leaving only jax.extend.core.primitives.jit_p
#30080 opened Jul 8, 2025 by copybara-service bot Loading…
@froystig
[Mosaic] Support multiple non-contracting dims if they are collapsable.
#30076 opened Jul 8, 2025 by copybara-service bot Loading…
Reverts e562cddc3421f3ee10a031ae21c48ed85f3edbb2
#30066 opened Jul 8, 2025 by copybara-service bot Loading…
[Mosaic GPU] Fix extract_constant_from_{least,most}_replicated_expression_for_hint type annotation.
#30063 opened Jul 8, 2025 by copybara-service bot Loading…
@bchetioui
On-the-fly trace-time DCE using Python ref counting kokoro:force-run pull ready Ready for copybara import and testing
#30062 opened Jul 8, 2025 by dougalm Loading…
Always lower ragged dot for cpu, gpu, and tpu
#30058 opened Jul 8, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add a layout attribute to mosaic_gpu.BroadcastInDimOp.
#30044 opened Jul 8, 2025 by copybara-service bot Loading…
@bchetioui
Support complex numbers in pallas uninitialised_value
#30043 opened Jul 8, 2025 by PhilipVinc Loading…
[Mosaic GPU] Add Reduce as a new constructor of Expressions.
#30042 opened Jul 8, 2025 by copybara-service bot Loading…
@bchetioui
[direct-linearize] fix tests to pass under JAX_USE_DIRECT_LINEARIZE=1 pull ready Ready for copybara import and testing
#30040 opened Jul 8, 2025 by mattjj Loading…
@mattjj
Bump hypothesis from 6.102.4 to 6.135.26 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#30027 opened Jul 7, 2025 by dependabot bot Loading…
@jakevdp
fix cudnn sdpa invalid seqlen for unused segments kokoro:force-run pull ready Ready for copybara import and testing
#30023 opened Jul 7, 2025 by Cjkkkk Loading…
@mattjj
1
Bump fonttools from 4.51.0 to 4.58.5 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#30019 opened Jul 7, 2025 by dependabot bot Loading…
ProTip! Filter pull requests by the default branch with base:main.