-
Notifications
You must be signed in to change notification settings - Fork 3.2k
Pull requests: jax-ml/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
[Pallas:MGPU][NFC] Consistently use
self.kernel instead of plgpu.kernel in tests.
#32405
by copybara-service
bot
was merged Oct 7, 2025
Loading…
[pallas:mgpu] Fix error message for debug_print in wg semantics.
#32401
by copybara-service
bot
was merged Oct 7, 2025
Loading…
Slightly tweak a test tolerance to avoid a test failure for gradients of acosh.
#32392
by copybara-service
bot
was merged Oct 7, 2025
Loading…
Attempting to turn on memories_tests for GPU again after fixes.
#32379
by copybara-service
bot
was merged Oct 6, 2025
Loading…
[Mosaic GPU][NFC] Simplify and document profiler logic.
#32372
by copybara-service
bot
was merged Oct 6, 2025
Loading…
[Mosaic TPU][relayout] Mask bitwidth change with any change of tiling, offset and implicit_dim
#32332
by copybara-service
bot
was merged Oct 6, 2025
Loading…
Remove deprecated non-hermetic CUDA and NCCL repo rules.
#32331
by copybara-service
bot
was merged Oct 3, 2025
Loading…
Relax tolerance for test_emit_with_no_output. When setting
--xla_gpu_experimental_enable_fusion_autotuner, this test will become flaky when it goes through the BlockLevelEmitter. It looks like the tolerance is too high because the differences seem quite small and the autotuner doesn't mark it as having incorrect results, but let me know if there's more I should check.
#32313
by copybara-service
bot
was merged Oct 6, 2025
Loading…
Clarify
optimization_barrier semantics.
#32309
by copybara-service
bot
was merged Oct 2, 2025
Loading…
In jax.scipy.linalg.block_diag, form a single wider concatenation rat…
kokoro:force-run
pull ready
Ready for copybara import and testing
#32302
by hawkinsp
was merged Oct 2, 2025
Loading…
[JAX] Adjusting mesh for allowing 2 cores per chip on V5P
#32298
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[Mosaic GPU] Add SMEM inference rules for
wgmma.
#32282
by copybara-service
bot
was merged Oct 3, 2025
Loading…
Update
rules_ml_toolchain to newer version with latest updates and SYCL/OneAPI support
#32273
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[JAX] Make colocated Python's serialization pickle each common object only once
#32271
by copybara-service
bot
was merged Oct 6, 2025
Loading…
Add support for GCS paths in jax.collect_profile
#32269
by copybara-service
bot
was merged Oct 6, 2025
Loading…
[JAX] Add a serialization roundtrip test for colocated Python
#32266
by copybara-service
bot
was closed Oct 1, 2025
Loading…
Speed up unpacking
.tar.xz files in the hermetic toolchain by using multithreaded xz execution.
#32262
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[Mosaic GPU] Add SMEM inference rules for
tcgen05.mma
#32258
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU] Add a trivial SMEM inference rule for
memref.view.
#32257
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU] Add an SMEM inference pass-through rule for
memref.cast.
#32256
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU] Properly handle SMEM in
scf.for.
#32255
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU] Implement SMEM inference rules for
mgpu.async_load/mgpu.async_store.
#32254
by copybara-service
bot
was closed Oct 3, 2025
Loading…
[Mosaic GPU] Add SMEM inference rules for
vector.load/vector.store.
#32253
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU][NFC] Specify the memory space for two existing tests.
#32250
by copybara-service
bot
was merged Oct 2, 2025
Loading…
Previous Next
ProTip!
Type g p on any issue or pull request to go back to the pull request listing page.