CI: 10/31/24 upstream sync#117

Closed

charleshofer wants to merge 10 commits intorocm-mainfrom main

+966-162

Commits on Oct 30, 2024

Re-land jax-ml#24589 with fixes to handle dtype that is not compatible with NumPy.

Jake VanderPlas
authored and
Google-ML-Automation
committed
Remove the variant since sparsecore is only on v5p and it's device kind is TPU v5.

yashk2810
authored and
Google-ML-Automation
committed

Commits on Oct 31, 2024

Remove UnshapedArray values from JAX (it remains as an abstract class).

dougalm
authored and
Google-ML-Automation
committed
[Mosaic] Extend tpu matmulop to have dimension dims. Add support for batching and simple transposition.
Google-ML-Automation
committed
[MOSAIC:GPU] Extend the mosaic mlir dialect with fragmented layouts.

dimitar-asenov
authored and
Google-ML-Automation
committed
[pallas:mosaic_gpu] plgpu.copy_smem_to_gmem no longer transparently commits SMEM

superbobry
authored and
Google-ML-Automation
committed
[Mosaic GPU] Add Python bindings for the Mosaic GPU MLIR dialect.

bchetioui
authored and
Google-ML-Automation
committed
Run dot algorithm tests with PJRT plugin.
dfm
committed
Merge pull request jax-ml#24480 from dfm:dot-algorithm-plugin-enable
Google-ML-Automation
committed
[Mosaic] Fix some imports.

18praveenb
authored and
Google-ML-Automation
committed