ROCm/jax

Pull Requests Commits

Enable scaled_matmul tests on ROCm

phambinhfin committed 72 days ago

6e91aac5

[lapack][forward_fix] Correctly suppress MSAN warnings from SVD when using a threadpool.

basioli-k committed 72 days ago

de358e76

Enable subset device execution for Mosaic GPU collectives with collective metadata.

Google-ML-Automation committed 72 days ago

5a72fb1c

Fix cross-assembly boundaries destruction of non-trivial objects

Google-ML-Automation committed 72 days ago

4bab92b3

[Pallas TPU] Add support for `jax.empty_ref`.

bchetioui committed 72 days ago

98d523a6

Support TMA initialization with collective metadata

Google-ML-Automation committed 72 days ago

b67da95b

[Mosaic GPU] Fix failures on older jaxlibs after MLIR upgrade.

bchetioui committed 72 days ago

d53d6c63

[mutable-arrays] Add `free_ref` primitive.

bchetioui committed 72 days ago

5c37cf48

Update XLA dependency to use revision http://github.com/openxla/xla/commit/301198e60a88c879842588263fac1eed5dbf3a04

Google-ML-Automation committed 72 days ago

00930e85

Make pull_block_spec only require the grid size instead of the full grid.

Davis Yoshida committed 73 days ago

cc989d14

Integrate LLVM at llvm/llvm-project@b4b8d4e5d062

Google-ML-Automation committed 73 days ago

6a3766a7

[Pallas] Add a `with_scoped` function decorator that is like `run_scoped` but does not run the function. This makes it more easily composable.

sharadmv committed 73 days ago

942d95c2

Bump rules_ml_toolchain to enable customization of CUDA Core Compute Libraries repository

yuriivcs committed 73 days ago

a7d3b223

Add a `TiledLayoutAttr::getContiguous` factory.

nvgrw committed 73 days ago

e3da4b93

Update XLA dependency to use revision http://github.com/openxla/xla/commit/7ab805eb9a772790dcb83f3fe8c561a895d7f23e

belitskiy committed 73 days ago

9b5f9626

Merge pull request #32115 from ROCm:jd/plugin_ver_upstream

Google-ML-Automation committed 73 days ago

a75f3831

Merge pull request #34809 from justinjfu:explicit_sharding

Google-ML-Automation committed 73 days ago

0bcf6e59

[XLA:MGPU] Port the entire TiledLayout class to C++.

golechwierowicz committed 73 days ago

accb7199

Update XLA dependency to use revision http://github.com/openxla/xla/commit/f9ede68ed63446045a0ab73660a80348c405b394

Google-ML-Automation committed 73 days ago

05d9eec3

Add reshard rule for mosaic gpu

justinjfu committed 73 days ago

a67c84af

Add JAX version metadata to xspace, creating a metadata registry API to add metadata

Google-ML-Automation committed 74 days ago

8ca2421b

Raise NotImplementedError for pcast primitives in their impl rule. `pcast` should be called under shard_map. Fixes https://github.com/jax-ml/jax/issues/34801

yashk2810 committed 74 days ago

3ddf0685

Merge pull request #34734 from jakevdp:indexing-dynamic-slice

Google-ML-Automation committed 74 days ago

c9ef1c10

[indexing] support pallas-like dynamic slicing in jax.numpy indexing

jakevdp committed 74 days ago

6156decd

[pallas:sc] `fetch_and_add` now support multiple indices

superbobry committed 74 days ago

b23e3a44

[pmap] Add reshard jet rule.

danielsuo committed 74 days ago

dc1eedbe

Increase shard count for the newly added scaled_dot_test on GPU and TPU.

belitskiy committed 74 days ago

c9dbcaa9

Merge pull request #34798 from jakevdp:fix-lint

Google-ML-Automation committed 74 days ago

e1055fdd

Fix lint at head

jakevdp committed 74 days ago

6f6141d2

[mosaic] Exported `Tiling` under `jax.experimental.mosaic`

superbobry committed 74 days ago

30d7f945

Older