ROCm/jax

Pull Requests Commits

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

AratiGanesh committed 18 days ago

Verified 18d0f073

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

Google-ML-Automation committed 19 days ago

1803bf2c

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

allanrenucci committed 19 days ago

fff793a0

[pallas:mosaic] Extracted `_lower_to_custom_call`

superbobry committed 19 days ago

d2aeeadf

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

Google-ML-Automation committed 19 days ago

73a9f058

Remove an extra assert from standard_abstract_eval.

yashk2810 committed 19 days ago

c2072648

Add a version guard around test_pspec_subclass_error

yashk2810 committed 20 days ago

d406251a

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

yashk2810 committed 20 days ago

438ed6d6

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

Google-ML-Automation committed 20 days ago

c516062a

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

Google-ML-Automation committed 20 days ago

1c7d68a4

Merge pull request #35088 from jax-ml:shmap_explicit

Google-ML-Automation committed 20 days ago

d5d63f46

shmap doc full explicit

yashk2810 committed 20 days ago

aebab97d

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

tlongeri committed 20 days ago

f6b0aa95

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

Google-ML-Automation committed 20 days ago

60bc8472

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

Google-ML-Automation committed 20 days ago

2254093b

[mosaic] Reverted `tpu_custom_call.lowered_as_tpu_kernel`

superbobry committed 21 days ago

6668ea68

Merge pull request #35082 from jakevdp:fix-jupytext

Google-ML-Automation committed 21 days ago

a18a7648

Fix jupytext lint error

jakevdp committed 21 days ago

77987113

[random] fix PRNGKeyArray pickle under Python 3.11

jakevdp committed 21 days ago

4712a6a9

Removed more dead code all over JAX

superbobry committed 21 days ago

b38a8164

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

Google-ML-Automation committed 21 days ago

019c6d63

[pallas:mosaic_gpu] Removed special-casing for `kernel_name` in `plgpu.kernel`

superbobry committed 21 days ago

bfaaac9b

[mosaic_gpu] Removed a few unused functions in `inference_utils`

superbobry committed 21 days ago

a6aab034

enable shmap docs build (again)

mattjj committed 21 days ago

8181e501

Merge pull request #35052 from jakevdp:map-bugs

Google-ML-Automation committed 21 days ago

2b2d52df

Update `rules_ml_toolchain` to pass CCCL headers to NVCC before all other headers for local CUDA.

ybaturina committed 21 days ago

04c74b37

[ROCm] Modified `test_with_memory_space` to include ROCm tests.

tsrw2048 committed 21 days ago

b68afcc4

[pallas] Updated CHANGELOG.md

superbobry committed 21 days ago

b6c489d9

Request a full h100 gpu when neccessary.

Google-ML-Automation committed 21 days ago

95ab165e

Merge pull request #35051 from jakevdp:pyrefly-fixes-4

Google-ML-Automation committed 21 days ago

fcf0aa55

Older