ROCm/jax

Pull Requests Commits

[ROCm] Enabled ROCm devices to default to Jacobi SVD on smaller matrices.

tsrw2048 committed 146 days ago

2dabef20

Prepare JAX test wrappers for rules_python 1.8.3 upgrade

ecalubaquib committed 146 days ago

ebd5249d

Merge pull request #34973 from jakevdp:more-pyrefly-fixes

Google-ML-Automation committed 147 days ago

c57f2524

Upgrade to CCCL 3.2.0 with better compatibility to CUDA 12.9, 13.0, and 13.1+.

ybaturina committed 147 days ago

3694b2be

[pallas] Pulled `default_memory_space` into `Mesh`

superbobry committed 147 days ago

d55be985

[typing] fix a number of pyrefly typechecking errors

jakevdp committed 147 days ago

002ac9bd

Merge pull request #34991 from superbobry:pyrefly

Google-ML-Automation committed 147 days ago

9f1399e5

Remove deprecated coordination service option.

mwhittaker committed 147 days ago

747397e1

Bump up shard counts for api_test, and scaled_dot_test, mostly for TSAN.

belitskiy committed 147 days ago

068ae53b

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

superbobry committed 147 days ago

316ffe98

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

bchetioui committed 147 days ago

41c6e2a8

[pallas:mosaic] Extracted `lower_jaxpr_into_module`

superbobry committed 147 days ago

e455dc10

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

golechwierowicz committed 147 days ago

4abfdc62

Bumped Ruff, mypy and Pyrefly versions in .pre-commit-config.yaml

superbobry committed 147 days ago

4cb6c61b

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

golechwierowicz committed 147 days ago

30427650

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

bchetioui committed 147 days ago

f9a10ff0

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

golechwierowicz committed 147 days ago

5dacd8a4

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

Google-ML-Automation committed 147 days ago

55513d6f

Merge pull request #34985 from gnecula:test_pr

Google-ML-Automation committed 147 days ago

238fdfab

[NFC] Remove unnecessary comment from `testMultivariateNormalSingularCovariance`.

bchetioui committed 147 days ago

3334de9a

Add an error message to assertValidVersion

gnecula committed 147 days ago

24c453ae

Merge pull request #34986 from gnecula:fix_test_error

Google-ML-Automation committed 147 days ago

0d3230b1

Merge pull request #34689 from ROCm:rocm_fix_pallas

Google-ML-Automation committed 147 days ago

ebbdfdce

Fix failing test in array_interoperability_test.py

gnecula committed 147 days ago

3166c321

Some fixes following #34911

superbobry committed 147 days ago

b8723498

Reverts 583744603265ea24dfaca6a1300e46b65377b3ef

h-joo committed 147 days ago

001c1ad2

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

Google-ML-Automation committed 147 days ago

6b3a3e9b

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

Google-ML-Automation committed 147 days ago

07ef410d

Automated Code Change

h-joo committed 147 days ago

58374460

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

Google-ML-Automation committed 147 days ago

1202f110

Older