ROCm/jax

Pull Requests Commits

Skip scaled matmul tests on CPU

phambinhfin committed 53 days ago

11558844

Address review comments

phambinhfin committed 63 days ago

731725ac

Enable neural network tests on ROCm

phambinhfin committed 70 days ago

a5426f8a

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

basioli-k committed 70 days ago

de358e76

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

Google-ML-Automation committed 70 days ago

5a72fb1c

Fix cross-assembly boundaries destruction of non-trivial objects

Google-ML-Automation committed 70 days ago

4bab92b3

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

bchetioui committed 70 days ago

98d523a6

Support TMA initialization with collective metadata

Google-ML-Automation committed 70 days ago

b67da95b

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

bchetioui committed 70 days ago

d53d6c63

[mutable-arrays] Add `free_ref` primitive.

bchetioui committed 70 days ago

5c37cf48

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

Google-ML-Automation committed 71 days ago

00930e85

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

Davis Yoshida committed 71 days ago

cc989d14

Integrate LLVM at llvm/llvm-project@b4b8d4e5d062

Google-ML-Automation committed 71 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 71 days ago

942d95c2

Bump rules_ml_toolchain to enable customization of CUDA Core Compute Libraries repository

yuriivcs committed 71 days ago

a7d3b223

Add a `TiledLayoutAttr::getContiguous` factory.

nvgrw committed 71 days ago

e3da4b93

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

belitskiy committed 71 days ago

9b5f9626

Merge pull request #32115 from ROCm:jd/plugin_ver_upstream

Google-ML-Automation committed 71 days ago

a75f3831

Merge pull request #34809 from justinjfu:explicit_sharding

Google-ML-Automation committed 71 days ago

0bcf6e59

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

golechwierowicz committed 72 days ago

accb7199

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

Google-ML-Automation committed 72 days ago

05d9eec3

Add reshard rule for mosaic gpu

justinjfu committed 72 days ago

a67c84af

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

Google-ML-Automation committed 72 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 72 days ago

3ddf0685

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

Google-ML-Automation committed 72 days ago

c9ef1c10

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

jakevdp committed 72 days ago

6156decd

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

superbobry committed 72 days ago

b23e3a44

[pmap] Add reshard jet rule.

danielsuo committed 72 days ago

dc1eedbe

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

belitskiy committed 72 days ago

c9dbcaa9

Merge pull request #34798 from jakevdp:fix-lint

Google-ML-Automation committed 72 days ago

e1055fdd

Older