Pull requests: jax-ml/jax
Pull requests list
Increasing shard count and removing asan builds to prevent timeouts.
#26175
by copybara-service
bot
was closed Jan 29, 2025
Loading…
Add win32 guard to fix imports on Windows
#26138
by copybara-service
bot
was merged Jan 28, 2025
Loading…
[Mosaic GPU] Support reads/writes from SMEM to WGMMARowFragLayout arrays.
#27166
by copybara-service
bot
was merged Mar 18, 2025
Loading…
[Mosaic GPU] Limit the maximum number of registers per thread to 255.
#27663
by copybara-service
bot
was merged Apr 4, 2025
Loading…
[Mosaic GPU] Support more layouts in the
swap lowering.
#27431
by copybara-service
bot
was merged Mar 26, 2025
Loading…
Updated the repr of GPU devices to be more consistent with TPUs/CPUs.
#22348
by copybara-service
bot
was merged Jul 11, 2024
Loading…
Use boolean values for partial mask blocks in the splash attention kernel.
#25323
by copybara-service
bot
was merged Dec 11, 2024
Loading…
Fix error message in dot_product_attention
#22951
by copybara-service
bot
was merged Aug 8, 2024
Loading…
Faster MHA backwards pass
pull ready
Ready for copybara import and testing
#22820
by Rifur13
was merged Aug 7, 2024
Loading…
[Mosaic GPU] Add lowering for
log, and a fast path using log2.
#26910
by copybara-service
bot
was merged Mar 4, 2025
Loading…
Increase shard count after adding more tests
#25411
by copybara-service
bot
was merged Dec 11, 2024
Loading…
Increase the absolute error tolerance to fix flaky tests.
#26413
by copybara-service
bot
was merged Feb 7, 2025
Loading…
[Mosaic GPU] Return the combined softmax residuals.
#27455
by copybara-service
bot
was merged Apr 4, 2025
Loading…
[Mosaic GPU] Adding pallas bindings to broadcast over the leading dimension and load a ref into WGMMAColFragLayout format.
#27356
by copybara-service
bot
was merged Mar 24, 2025
Loading…
[Mosaic GPU] Adding a new layout WGMMAColFragLayout to be able to load a 1d array and broadcast it along the leading dimension to a 2d shape as an input to a wgmma.
#27276
by copybara-service
bot
was merged Mar 24, 2025
Loading…
[Mosaic GPU] Adding a primitive to load from memrefs *with* a specified layout.
#27207
by copybara-service
bot
was merged Mar 24, 2025
Loading…
[Mosaic GPU] Reverse the direction in which we iterate through the pallas grid to be row-major.
#29105
by copybara-service
bot
was merged Jul 25, 2025
Loading…
[Pallas][Mosaic GPU] Fix racy test by having each warpgroup work on non-overlapping data.
#30354
by copybara-service
bot
was merged Jul 21, 2025
Loading…
[Mosaic GPU] Adding a deterministic backwards pass to the pallas MGPU kernel. It's implemented as split dq and dkv kernels to have enough SMEM for double-buffering.
#27860
by copybara-service
bot
was merged Apr 24, 2025
Loading…
[Pallas][Mosaic GPU] Support pytree in/out specs in the warp specialized pipeline.
#30414
by copybara-service
bot
was merged Aug 1, 2025
Loading…
[Mosaic GPU] Increasing number of shards because tests have been timing out in asan mode after adding https://github.com/jax-ml/jax/pull/27860.
#28354
by copybara-service
bot
was closed May 9, 2025
Loading…
[Mosaic GPU] Adding an optional causal masking to the manual pipelining example.
#28508
by copybara-service
bot
was merged May 19, 2025
Loading…
ProTip!
Type g p on any issue or pull request to go back to the pull request listing page.