ROCm/jax

Pull Requests Commits

[typing] fix some pyrefly unbound local errors

jakevdp committed 11 days ago

beebc8ed

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

Google-ML-Automation committed 11 days ago

2ce249ed

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

Google-ML-Automation committed 11 days ago

a1dea76f

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

Google-ML-Automation committed 11 days ago

f404b9af

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

Google-ML-Automation committed 11 days ago

336e5872

Merge pull request #34494 from ROCm:rocm_gpu_lowering_gesvdj

Google-ML-Automation committed 11 days ago

a5a62d6f

Add an extra cuBLAS version check.

bchetioui committed 11 days ago

dc150667

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

superbobry committed 11 days ago

0f3ac445

Merge pull request #35025 from gnecula:pl_cdiv

Google-ML-Automation committed 11 days ago

b1d0a5bd

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

bchetioui committed 11 days ago

90b71283

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

gnecula committed 11 days ago

15dedfe9

Merge pull request #35020 from gnecula:fix_typo

Google-ML-Automation committed 11 days ago

8e5a0fe5

Removed a few more unused functions detected by Vulture

superbobry committed 11 days ago

6d6d1784

Fix typo and line lengths

gnecula committed 11 days ago

958323f1

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

Google-ML-Automation committed 11 days ago

616e7b79

Merge pull request #35018 from gnecula:empty_pr

Google-ML-Automation committed 11 days ago

52cfe894

Fixed the linter

gnecula committed 11 days ago

62d2506e

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

tsrw2048 committed 11 days ago

17337fef

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

Google-ML-Automation committed 11 days ago

c8c228c6

[Mosaic:GPU] Use collective kernel API directly.

Google-ML-Automation committed 11 days ago

8949bce9

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

tsrw2048 committed 11 days ago

72c052e8

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

Google-ML-Automation committed 11 days ago

5479321e

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

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

991fb76b

Merge pull request #34977 from mattjj:hijax-vmappable

Google-ML-Automation committed 11 days ago

ce17815f

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

mattjj committed 12 days ago

eab55a49

Remove callable out_specs support from internal `_shard_map`. It existed of pmap_shmap merge but it doesn't use the callable support anymore

yashk2810 committed 12 days ago

7ff6c399

Add ROCm backward compatibility test for threefry2x32

AratiGanesh committed 12 days ago

bc4aa31b

Add ROCm support to annotate_device_placement backward compatibility test

AratiGanesh committed 12 days ago

5cc2dd8a

[ROCm] Enabled the Triton Pallas tests to run for ROCm.

tsrw2048 committed 12 days ago

a1423f34

Newer Older