ROCm/jax

Pull Requests Commits

Remove incorrect ROCm lowering for scaled_matmul to prevent segfault

Ruturaj4 committed 36 days ago

4d210c5d

[pallas] Deduplicated abstract_eval and typecheck rules for `core_map_p`

superbobry committed 37 days ago

316ffe98

[Mosaic GPU][NFC] Minor code simplification in `mosaic_gpu_ext.cc`.

bchetioui committed 37 days ago

41c6e2a8

[pallas:mosaic] Extracted `lower_jaxpr_into_module`

superbobry committed 37 days ago

e455dc10

[XLA:MGPU] Port TileTransform to C++.

golechwierowicz committed 37 days ago

4abfdc62

[XLA:MGPU] Move more functions to utils and make NdIndices accept callback.

golechwierowicz committed 37 days ago

30427650

[Mosaic GPU] Fix bugs in `transform_type` implementation.

bchetioui committed 37 days ago

f9a10ff0

[XLA:MGPU] Add a convenience 'c' function constructing ConstantOp

golechwierowicz committed 37 days ago

5dacd8a4

[Pallas/Mosaic GPU] Use new API for collective metadata creation

Google-ML-Automation committed 37 days ago

55513d6f

Merge pull request #34985 from gnecula:test_pr

Google-ML-Automation committed 37 days ago

238fdfab

[NFC] Remove unnecessary comment from `testMultivariateNormalSingularCovariance`.

bchetioui committed 37 days ago

3334de9a

Add an error message to assertValidVersion

gnecula committed 37 days ago

24c453ae

Merge pull request #34986 from gnecula:fix_test_error

Google-ML-Automation committed 37 days ago

0d3230b1

Merge pull request #34689 from ROCm:rocm_fix_pallas

Google-ML-Automation committed 37 days ago

ebbdfdce

Fix failing test in array_interoperability_test.py

gnecula committed 37 days ago

3166c321

Some fixes following #34911

superbobry committed 37 days ago

b8723498

Reverts 583744603265ea24dfaca6a1300e46b65377b3ef

h-joo committed 37 days ago

001c1ad2

Merge pull request #34578 from ROCm:ups-select-skip-conv

Google-ML-Automation committed 37 days ago

6b3a3e9b

Merge pull request #34534 from magaonka-amd:enable-lax-dot-rocm-jax-ml

Google-ML-Automation committed 37 days ago

07ef410d

Automated Code Change

h-joo committed 37 days ago

58374460

Merge pull request #34577 from ROCm:rocm-enable-workspace-overflow-test

Google-ML-Automation committed 37 days ago

1202f110

Merge pull request #34567 from ROCm:unskip_multivariate_normal_rocm

Google-ML-Automation committed 37 days ago

fb31d9fd

Merge pull request #34571 from magaonka-amd:enable-reduce-window-tests-rocm-upstream

Google-ML-Automation committed 37 days ago

4c157782

Merge pull request #34890 from magaonka-amd:rocm-enable-memory-space-export-tests

Google-ML-Automation committed 37 days ago

14a1491e

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

Google-ML-Automation committed 37 days ago

8d9e1809

Use `reshape` in `_reshape_transpose_rule` instead of `reshape_p.bind` as it was before unreduced/reduced support was added to reshape_p.

yashk2810 committed 37 days ago

623d9dba

Merge pull request #34304 from ROCm:rocm_psend_support

Google-ML-Automation committed 37 days ago

d76e2431

Merge pull request #34662 from magaonka-amd:enable-caitojax-rocm

Google-ML-Automation committed 37 days ago

ef18ec70

Merge pull request #34470 from ROCm:rocm_enable_tridiagonal_tests

Google-ML-Automation committed 37 days ago

dc51c478

Merge pull request #34959 from ROCm:skip-tridiagonal-solve-grad-upstream-main

Google-ML-Automation committed 37 days ago

7136c7ab

Older