Skip to content

Pull requests: jax-ml/jax

New pull request New
630 Open 20,256 Closed
630 Open 20,256 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

Count jit_cpp_cache_misses by registering an event in def _pjit_call_impl_python
#30529 opened Jul 26, 2025 by copybara-service bot Loading…
@yashk2810
[pallas] Add support for fusing into a fusible within a vmap.
#30527 opened Jul 26, 2025 by copybara-service bot Loading…
@chr1sj0nes
Fix ASAN build.
#30524 opened Jul 26, 2025 by copybara-service bot Loading…
[pallas:mgpu] Switch batch dim position in kernel vmap rule to match new row-major ordering.
#30522 opened Jul 25, 2025 by copybara-service bot Loading…
@chr1sj0nes
Fix jaxlib ASAN github test
#30521 opened Jul 25, 2025 by copybara-service bot Loading…
Reorganize the tutorials
#30519 opened Jul 25, 2025 by melissawm Loading…
@emilyfertig
3
fix bnth layout shardy and infer_fwd_ouput_sharding
#30516 opened Jul 25, 2025 by Cjkkkk Loading…
Only bind a reduce_sum_p if it has nontrivial axes
#30515 opened Jul 25, 2025 by copybara-service bot Loading…
@sharadmv
Add support for scipy.special.erfcx (scaled complementary error funct…
#30506 opened Jul 25, 2025 by zk-quantum Loading…
1
[pallas] Deprecated pl.load and pl.store
#30505 opened Jul 25, 2025 by copybara-service bot Loading…
@superbobry
Fix OSS breakage due to TPU test.
#30503 opened Jul 25, 2025 by copybara-service bot Loading…
[Mosaic GPU] Make the vector length configurable in the TMEM_NATIVE_LAYOUT
#30501 opened Jul 25, 2025 by copybara-service bot Loading…
@apaszke
Add fuser.py in jax wheel.
#30500 opened Jul 25, 2025 by copybara-service bot Loading…
[Mosaic GPU] Support wargroup semantic for automatic TMEM allocation.
#30499 opened Jul 25, 2025 by copybara-service bot Loading…
Internal changes.
#30494 opened Jul 25, 2025 by copybara-service bot Loading…
@danielsuo
Add Bazel TPU tests in JAX presubmit and postsubmit.
#30492 opened Jul 25, 2025 by copybara-service bot Loading…
Pallas/Mosaic support host output.
#30489 opened Jul 24, 2025 by copybara-service bot Loading…
Pipe incarnations through C plugin.
#30488 opened Jul 24, 2025 by copybara-service bot Loading…
@mwhittaker
core.Tracer shouldn't inherit from basearray.Array kokoro:force-run pull ready Ready for copybara import and testing
#30487 opened Jul 24, 2025 by mattjj Loading…
@mattjj
[XLA:GPU] update deadlocks verifier behavior to get computations from the scheduler
#30486 opened Jul 24, 2025 by copybara-service bot Loading…
fix: typo in direct_linearize_migration.md
#30485 opened Jul 24, 2025 by pfackeldey Loading…
[Mosaic GPU] Implement lowering rules for async_tmem_load and async_tmem_store.
#30484 opened Jul 24, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add async_tmem_load and async_tmem_store ops to the Mosaic MLIR dialect.
#30483 opened Jul 24, 2025 by copybara-service bot Loading…
[Mosaic GPU] Predicate tmem_alloc and tmem_dealloc at the warp level.
#30482 opened Jul 24, 2025 by copybara-service bot Loading…
Enable libnvptxcompiler and libnvjitlink support in OSS by default.
#30480 opened Jul 24, 2025 by copybara-service bot Loading…
ProTip! Filter pull requests by the default branch with base:main.