ROCm/jax

Pull Requests Commits

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

Google-ML-Automation committed 26 days ago

48ee0b1f

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

Google-ML-Automation committed 26 days ago

60a5ecb8

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

basioli-k committed 26 days ago

de358e76

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

Google-ML-Automation committed 26 days ago

5a72fb1c

Trigger CI/CD pipelinee

alekstheod committed 26 days ago

0c856ce1

Fix cross-assembly boundaries destruction of non-trivial objects

Google-ML-Automation committed 26 days ago

4bab92b3

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

bchetioui committed 26 days ago

98d523a6

Support TMA initialization with collective metadata

Google-ML-Automation committed 26 days ago

b67da95b

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

bchetioui committed 26 days ago

d53d6c63

[mutable-arrays] Add `free_ref` primitive.

bchetioui committed 26 days ago

5c37cf48

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

Google-ML-Automation committed 26 days ago

00930e85

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

Davis Yoshida committed 26 days ago

cc989d14

Integrate LLVM at llvm/llvm-project@b4b8d4e5d062

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

942d95c2

Bump rules_ml_toolchain to enable customization of CUDA Core Compute Libraries repository

yuriivcs committed 26 days ago

a7d3b223

Add a `TiledLayoutAttr::getContiguous` factory.

nvgrw committed 26 days ago

e3da4b93

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

belitskiy committed 26 days ago

9b5f9626

Address review comments

alekstheod committed 26 days ago

d21547af

Standardize parameter naming across workflows and actions

mminutoli committed 26 days ago

675e05cb

Adress review comments

alekstheod committed 26 days ago

d151959c

refactor: replace assert False with proper ValueError in serialization

hrideymarwah15 committed 26 days ago

f069eade

alekstheod committed 27 days ago

4441a183

[indexing] refactor to_static_slice & to_dynamic_slice

jakevdp committed 27 days ago

fc2dba41

alekstheod committed 27 days ago

d583f1a0

Merge pull request #32115 from ROCm:jd/plugin_ver_upstream

Google-ML-Automation committed 27 days ago

a75f3831

Address review comments, use repo name as a convention

alekstheod committed 27 days ago

a58e5e5d

Address review comments

alekstheod committed 27 days ago

40462512

Address review comments, do not use env variables

alekstheod committed 27 days ago

6cb71238

Use Label() instead

alekstheod committed 27 days ago

b2be9cca

Merge pull request #34809 from justinjfu:explicit_sharding

Google-ML-Automation committed 27 days ago

0bcf6e59

Newer Older