ROCm/jax

Pull Requests Commits

Add c128 test data for ROCm cholesky solver backward compatibility test

AratiGanesh committed 106 days ago

a2a7f688

Add ROCm backward compatibility test for cholesky solver

AratiGanesh committed 106 days ago

536bb3f8

Merge pull request #34474 from ROCm:rocm_gesvdj_support

Google-ML-Automation committed 106 days ago

bad6b6ea

Merge pull request #35010 from jakevdp:pyrefly-fixes-3

Google-ML-Automation committed 106 days ago

38d3d87e

[typing] fix some pyrefly unbound local errors

jakevdp committed 106 days ago

beebc8ed

Merge pull request #34802 from magaonka-amd:enable-sparse-tests-rocm-upstream

Google-ML-Automation committed 106 days ago

2ce249ed

Merge pull request #34774 from magaonka-amd:enable-lobpcg-tests-rocm

Google-ML-Automation committed 106 days ago

a1dea76f

Merge pull request #34575 from ROCm:rocm-enable-triton-pallas-tests

Google-ML-Automation committed 106 days ago

f404b9af

Merge pull request #34893 from magaonka-amd:rocm-fix-aot-test-platform

Google-ML-Automation committed 106 days ago

336e5872

Merge pull request #34494 from ROCm:rocm_gpu_lowering_gesvdj

Google-ML-Automation committed 106 days ago

a5a62d6f

Add an extra cuBLAS version check.

bchetioui committed 106 days ago

dc150667

[pallas:sc] Slightly tweaked `PipelineTest.test_basic` under 7x

superbobry committed 106 days ago

0f3ac445

Merge pull request #35025 from gnecula:pl_cdiv

Google-ML-Automation committed 106 days ago

b1d0a5bd

[Mosaic GPU] Implement `tile_shape`, a MLIR operation that allows tiling a memref.

bchetioui committed 106 days ago

90b71283

[pallas] Generalize pl.cdiv to work with symbolic dimensions

gnecula committed 106 days ago

15dedfe9

Merge pull request #35020 from gnecula:fix_typo

Google-ML-Automation committed 106 days ago

8e5a0fe5

Removed a few more unused functions detected by Vulture

superbobry committed 106 days ago

6d6d1784

Fix typo and line lengths

gnecula committed 106 days ago

958323f1

Merge pull request #34875 from ROCm:rocm-threefry2x32-test

Google-ML-Automation committed 106 days ago

616e7b79

Merge pull request #35018 from gnecula:empty_pr

Google-ML-Automation committed 106 days ago

52cfe894

Fixed the linter

gnecula committed 106 days ago

62d2506e

PR #34822: [ROCm] Added back compatibility test for hipsolver_sytrd

tsrw2048 committed 106 days ago

17337fef

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

Google-ML-Automation committed 106 days ago

c8c228c6

[Mosaic:GPU] Use collective kernel API directly.

Google-ML-Automation committed 106 days ago

8949bce9

PR #34862: [ROCm] Added back compatibility test for hipsolver_geqrf

tsrw2048 committed 106 days ago

72c052e8

Merge pull request #34832 from ROCm:add-rocm-device-placement-compat-test

Google-ML-Automation committed 106 days ago

5479321e

Merge pull request #34869 from ROCm:rocm-gesvd-compat-test

Google-ML-Automation committed 106 days ago

7a197166

[Take 2] In explicit mode, error out by default if the specs of arguments do not match in_specs passed to shard_map. In other words, in_specs in explicit_mode acts as an assert to make sure (`jax.typeof(arg).sharding.spec == in_specs`)

yashk2810 committed 106 days ago

991fb76b

Merge pull request #34977 from mattjj:hijax-vmappable

Google-ML-Automation committed 106 days ago

ce17815f

[hijax] add HiType.{inc_rank,dec_rank} for vmap and scan.

mattjj committed 106 days ago

eab55a49

Older