Pull requests: jax-ml/jax
Pull requests list
[Mosaic GPU] Disable
mypy in layout inference files.
#32333
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Mosaic GPU][NFC] Clean up type checking errors.
#32324
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Pallas/Mosaic GPU] Extend test coverage for
hopper_mixed_type_matmul_mgpu.
#32329
by copybara-service
bot
was closed Oct 3, 2025
Loading…
[Mosaic GPU] Signal pass failure in
AssemblyToBinaryPass if PTX compilation failed.
#32330
by copybara-service
bot
was merged Oct 3, 2025
Loading…
[Pallas][Mosaic GPU] Improve ragged dot kernel for blackwell.
#31708
by copybara-service
bot
was merged Sep 15, 2025
Loading…
[Mosaic TPU] Use integer types for transpose emulation
#32289
by copybara-service
bot
was merged Oct 3, 2025
Loading…
Fix out_memory_spaces_flat calculation. It was using in_tree before but should have used out_tree instead.
#32322
by copybara-service
bot
was merged Oct 3, 2025
Loading…
Add experimental higher-order primitive called compute_on2. This should eventually replace the current
compute_on.
#32319
by copybara-service
bot
was merged Oct 3, 2025
Loading…
Support collective cancelling in TFRT GPU client.
#30729
by copybara-service
bot
was closed Oct 2, 2025
Loading…
[pmap] Fix up remaining docstring examples and tests.
#32306
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[sparse_bcoo_bcsr] Reduce the number of sampled tests for cpu variants to avoid timeouts.
#32304
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[Pallas/Fuser] Add basic hitype support to fuser.fuse
#32277
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[Pallas/Mosaic GPU][Docs] Add missing backtick to reference.
#32301
by copybara-service
bot
was merged Oct 2, 2025
Loading…
[pallas:sc]
pltpu.async_copy can now do dynamically-sized scatters and gathers
#32288
by copybara-service
bot
was merged Oct 2, 2025
Loading…
ProTip!
Find all pull requests that aren't related to any open issues with -linked:issue.