Skip to content

Pull requests: jax-ml/jax

New pull request New
554 Open 22,697 Closed
554 Open 22,697 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 deactivation of array_test on certain CPU platforms.
#33236 opened Nov 11, 2025 by copybara-service bot Loading…
@emilyfertig
Automated Code Change
#33234 opened Nov 11, 2025 by copybara-service bot Loading…
Fix mishandling of large matrices in batched eigendecomposition.
#33233 opened Nov 11, 2025 by copybara-service bot Loading…
1
@hawkinsp
Aot
#33225 opened Nov 10, 2025 by danielsuo Loading…
7
[Mosaic] Determine matmul data format from LHS instead of RHS.
#33223 opened Nov 10, 2025 by copybara-service bot Loading…
[remat] Make prevent_cse's tuple form account for consts
#33217 opened Nov 10, 2025 by gspschmid Loading…
2
Support pytree classes with dunder-style flatten/unflatten methods
#33215 opened Nov 10, 2025 by jakevdp Loading…
1
@jakevdp
5
Add wheel import tests for JAX and jaxlib.
#33214 opened Nov 10, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add support for 8-bit types in sparse tcgen05 MMA
#33213 opened Nov 10, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add support for integer WGMMA with LHS in registers under WG semantic.
#33212 opened Nov 10, 2025 by copybara-service bot Loading…
[Mosaic GPU] Grow k_group_elems instead of complaining about small swizzle
#33209 opened Nov 10, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Support nvfp4 scaled MMA instructions
#33208 opened Nov 10, 2025 by copybara-service bot Loading…
@apaszke
Add support for global scale to scaled_matmul_wrapper
#33204 opened Nov 10, 2025 by sergey-kozub Loading…
@justinjfu
5
[Mosaic GPU] Add custom vector load/store ops to the dialect.
#33202 opened Nov 10, 2025 by copybara-service bot Loading…
Automated Code Change
#33195 opened Nov 9, 2025 by copybara-service bot Loading…
Automated Code Change
#33194 opened Nov 9, 2025 by copybara-service bot Loading…
Fix ROCm lib names copy pattern to use wildcard version
#33187 opened Nov 7, 2025 by phambinhfin Loading…
3
Make nvidia version data optional for ROCm builds
#33186 opened Nov 7, 2025 by phambinhfin Loading…
@hawkinsp
4
Skip test_oom_compilation_only on TPU v6e and v7x
#33185 opened Nov 7, 2025 by copybara-service bot Loading…
@kanglant
Sici Asymptotic Expansion
#33178 opened Nov 7, 2025 by michaeleliot Loading…
1
@jakevdp
13
[Mosaic GPU] Add support for block-scaled MMA with all K sizes divisible by 128
#33173 opened Nov 7, 2025 by copybara-service bot Loading…
@apaszke
Add a test for a compiler bug related to int div/mod.
#33170 opened Nov 7, 2025 by copybara-service bot Loading…
Automated Code Change
#33169 opened Nov 7, 2025 by copybara-service bot Loading…
[Mosaic TPU][NFC] Share checks between the reshape/store and load/reshape optimization rules
#33163 opened Nov 7, 2025 by copybara-service bot Loading…
@apaszke
Fix Python path for executing "hi" JAX code.
#33161 opened Nov 7, 2025 by copybara-service bot Loading…
@chr1sj0nes
ProTip! no:milestone will show everything without a milestone.