ROCm/jax

Pull Requests Commits

Merge branch 'main' into tot-rocm-skips

gulsumgudukbay committed 6 days ago

Verified d83c6db3

[pallas] Reduced the number of usages of `JAX_PALLAS_USE_MOSAIC_GPU` in tests

superbobry committed 6 days ago

8f030ce6

Removed unused `format_test_name_suffix`

superbobry committed 6 days ago

49b604b7

[pallas] `pallas_call` no longer accepts `backend=`

superbobry committed 6 days ago

62816fe5

Removed a few more unused APIs

superbobry committed 6 days ago

218b59a0

Merge pull request #34829 from ROCm:add-rocm-lu-backward-compat-test

Google-ML-Automation committed 6 days ago

b60f218b

Don't assume Pallas:Triton is available (it's not on Windows)

apaszke committed 6 days ago

d64deca3

apaszke committed 6 days ago

a8bb1bce

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

Google-ML-Automation committed 7 days ago

c1db577d

Merge branch 'main' into add-rocm-lu-backward-compat-test

AratiGanesh committed 7 days ago

Verified 18d0f073

Merge pull request #34929 from ROCm:rocm-serial-compat-test

Google-ML-Automation committed 7 days ago

1803bf2c

[Mosaic GPU][NFC] Use newly introduced `ffi::Ffi::BindExecute()`.

allanrenucci committed 7 days ago

fff793a0

[pallas:mosaic] Extracted `_lower_to_custom_call`

superbobry committed 7 days ago

d2aeeadf

Update XLA dependency to use revision http://github.com/openxla/xla/commit/3eb5197b71e3771a69defc25e491b1299d6b93f3

Google-ML-Automation committed 8 days ago

73a9f058

Remove an extra assert from standard_abstract_eval.

yashk2810 committed 8 days ago

c2072648

Add a version guard around test_pspec_subclass_error

yashk2810 committed 8 days ago

d406251a

Mark `PartitionSpec` as `is_final()` so that it can't be subclassed.

yashk2810 committed 8 days ago

438ed6d6

Update XLA dependency to use revision http://github.com/openxla/xla/commit/1250c33dd0f79688c90f6ef1adb8ee333d9cae68

Google-ML-Automation committed 9 days ago

c516062a

Merge pull request #34966 from ROCm:skip-jvp13-upstream-main

Google-ML-Automation committed 9 days ago

1c7d68a4

Merge pull request #35088 from jax-ml:shmap_explicit

Google-ML-Automation committed 9 days ago

d5d63f46

shmap doc full explicit

yashk2810 committed 9 days ago

aebab97d

[Mosaic:TPU] Fold tpu.erase_layout into tpu.load and tpu.store

tlongeri committed 9 days ago

f6b0aa95

Merge pull request #34561 from ROCm:rocm-enable-cond-toeplitz-tests

Google-ML-Automation committed 9 days ago

60bc8472

Merge pull request #35081 from jakevdp:fix-prng-pickle

Google-ML-Automation committed 9 days ago

2254093b

[mosaic] Reverted `tpu_custom_call.lowered_as_tpu_kernel`

superbobry committed 9 days ago

6668ea68

Merge pull request #35082 from jakevdp:fix-jupytext

Google-ML-Automation committed 9 days ago

a18a7648

Fix jupytext lint error

jakevdp committed 9 days ago

77987113

[random] fix PRNGKeyArray pickle under Python 3.11

jakevdp committed 9 days ago

4712a6a9

Removed more dead code all over JAX

superbobry committed 9 days ago

b38a8164

Merge pull request #35047 from mattjj:shmap-docs-again

Google-ML-Automation committed 9 days ago

019c6d63

Older