ROCm/jax

Pull Requests Commits

[ROCm] Added back compatibility test for GESVD/GESVDJ on ROCm devices.

tsrw2048 committed 2 days ago

c08c4abe

[Mosaic:TPU] Make tpu.load and tpu.store's sublane_stride a DefaultValuedAttr instead of an OptionalAttr

tlongeri committed 2 days ago

36f475ec

Raise an error if indivisible shardings is returned by XLA.

yashk2810 committed 2 days ago

822e9a7c

Merge pull request #34641 from ROCm:add_support_to_dynamic_set_of_test_deps

Google-ML-Automation committed 3 days ago

9b90ff2a

Merge pull request #34450 from ROCm:upstream_rocm_ci

Google-ML-Automation committed 3 days ago

1627be1e

[Pallas/interpreter] Extend parameterized test cases to cover up to 16 Pallas threads.

Google-ML-Automation committed 3 days ago

48ee0b1f

Remove the check since the same thunk can be reinitialized in a while loop.

Google-ML-Automation committed 3 days ago

60a5ecb8

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

basioli-k committed 3 days ago

de358e76

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

Google-ML-Automation committed 3 days ago

5a72fb1c

Trigger CI/CD pipelinee

alekstheod committed 3 days ago

0c856ce1

Fix cross-assembly boundaries destruction of non-trivial objects

Google-ML-Automation committed 3 days ago

4bab92b3

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

bchetioui committed 3 days ago

98d523a6

Support TMA initialization with collective metadata

Google-ML-Automation committed 3 days ago

b67da95b

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

bchetioui committed 3 days ago

d53d6c63

[mutable-arrays] Add `free_ref` primitive.

bchetioui committed 3 days ago

5c37cf48

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

Google-ML-Automation committed 3 days ago

00930e85

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

Davis Yoshida committed 3 days ago

cc989d14

Integrate LLVM at llvm/llvm-project@b4b8d4e5d062

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

942d95c2

Bump rules_ml_toolchain to enable customization of CUDA Core Compute Libraries repository

yuriivcs committed 3 days ago

a7d3b223

Add a `TiledLayoutAttr::getContiguous` factory.

nvgrw committed 3 days ago

e3da4b93

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

belitskiy committed 3 days ago

9b5f9626

Address review comments

alekstheod committed 3 days ago

d21547af

Standardize parameter naming across workflows and actions

mminutoli committed 3 days ago

675e05cb

Adress review comments

alekstheod committed 3 days ago

d151959c

alekstheod committed 3 days ago

4441a183

alekstheod committed 3 days ago

d583f1a0

Merge pull request #32115 from ROCm:jd/plugin_ver_upstream

Google-ML-Automation committed 3 days ago

a75f3831

Address review comments, use repo name as a convention

alekstheod committed 3 days ago

a58e5e5d

Address review comments

alekstheod committed 4 days ago

40462512

Older