Skip to content

Pull requests: jax-ml/jax

New pull request New
425 Open 17,019 Closed
425 Open 17,019 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
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Pallas/Mosaic GPU] Add initial support for warpgroup semantics in lowering.
#26390 opened Feb 7, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic TPU] Fold VectorPack and VectorUnpack.
#26389 opened Feb 7, 2025 by copybara-service bot Loading…
[Mosaic TPU] Support bf16 div if HW does not directly support.
#26388 opened Feb 7, 2025 by copybara-service bot Loading…
Add sharding support for dynamic masks in TPU splash attention
#26387 opened Feb 7, 2025 by Rifur13 Loading…
[better_errors] Continue adding debug info to Jaxprs (step 5) pull ready Ready for copybara import and testing
#26386 opened Feb 7, 2025 by gnecula Loading…
@gnecula
[better_errors] Continue adding debug info to Jaxprs (step 4) pull ready Ready for copybara import and testing
#26384 opened Feb 7, 2025 by gnecula Loading…
@gnecula
refactor: move sorting ops out of lax_numpy pull ready Ready for copybara import and testing
#26383 opened Feb 7, 2025 by jakevdp Loading…
@jakevdp
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
#26382 opened Feb 7, 2025 by copybara-service bot Loading…
@dfm
Create JAX wheel build target.
#26378 opened Feb 7, 2025 by copybara-service bot Loading…
[mutable-arrays] persist shardings through xla computations
#26377 opened Feb 7, 2025 by mattjj Loading…
2
refactor: move array creation routines out of lax_numpy.py pull ready Ready for copybara import and testing
#26376 opened Feb 6, 2025 by jakevdp Loading…
@jakevdp
Expose checkpoint name in cuDNN SDPA
#26374 opened Feb 6, 2025 by Cjkkkk Loading…
1
Part 1 of a new autodidax based on "stackless"
#26373 opened Feb 6, 2025 by dougalm Loading…
Match output container to result_shape_dtypes in ffi_call.
#26371 opened Feb 6, 2025 by copybara-service bot Loading…
@dfm
Fix invalid lowerings for ROCm in Pallas
#26368 opened Feb 6, 2025 by mrodden Loading…
POC for GPU e2e feature testing
#26367 opened Feb 6, 2025 by sfvaroglu • Draft
[Mosaic GPU] Add support for 2-CTA MMA on Blackwell
#26364 opened Feb 6, 2025 by copybara-service bot Loading…
@apaszke
[pallas:mgpu] Fix and test multiple indexers where one is a dynamic selection index.
#26362 opened Feb 6, 2025 by copybara-service bot Loading…
[Mosaic GPU] Handle TMEM allocation in the compiler
#26360 opened Feb 6, 2025 by copybara-service bot Loading…
@apaszke
[pallas:triton] Lift dot_general restriction on minimal tile size for a.
#26355 opened Feb 6, 2025 by copybara-service bot Loading…
Add type annotation for zero_from_primal
#26346 opened Feb 6, 2025 by NeilGirdhar Loading…
Scaled matmul for mxfp8
#26345 opened Feb 5, 2025 by wenscarl • Draft
disable cuDNN flash attention head dim 256 on blackwell pull ready Ready for copybara import and testing
#26344 opened Feb 5, 2025 by Cjkkkk Loading…
3
[pallas:mosaic_gpu] Fixed a bug in how delay_release is handled in emit_pipeline
#26337 opened Feb 5, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Add tests for various tcgen05.mma configurations
#26335 opened Feb 5, 2025 by copybara-service bot Loading…
@apaszke
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.