ROCm/jax

Pull Requests Commits

Add possible mitigation to the crashes.

mminutoli committed 118 days ago

bfa68e59

Using --dist=loadscope to minimize the damage from crashes.

mminutoli committed 118 days ago

9175652f

Remove the limit on the number of failures in pytest command.

mminutoli committed 118 days ago

f4aa63f2

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 118 days ago

0ba4a780

Add download-jax-rocm-wheels composite action

mminutoli committed 118 days ago

ee289355

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 118 days ago

ac116a5a

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 118 days ago

9e737f71

adding run_pytest_rocm.sh

gulsumgudukbay committed 119 days ago

6421c5ea

Reverts f44bc6e37f770a9d79d24061f7960235a2892407

jakeharmon8 committed 209 days ago

403977d2

Prepare for JAX release 0.8.0

jakeharmon8 committed 209 days ago

b1ce4c57

[Pallas] Fix 64-bit multiplication emulation in Philox kernel

justinjfu committed 209 days ago

07dad8e8

Update XLA hash for JAX release

jakeharmon8 committed 209 days ago

4f06e7a6

Fix race when dialects were registered in an MLIR context when it was being accessed concurrently.

hawkinsp committed 209 days ago

1bd3402c

Clean up dead code.

brianwa84 committed 209 days ago

b942e894

Remove cusolver_potrf_ffi from the list of covered backwards compatibility targets.

gnecula committed 209 days ago

ac87cbe0

[Pallas:MGPU] Expose multicast TMA stores

apaszke committed 209 days ago

69fd2052

[Pallas:MGPU] Add an all-gather implementation

apaszke committed 209 days ago

8eed7dc5

Add tpu_sc.subcore_barrier primitive.

brianwa84 committed 209 days ago

6aae6f43

Add support for -1 size to ref reshaper, for convenience.

brianwa84 committed 209 days ago

8674d771

[Pallas:MGPU] Add ReducedLayout support to inline_mgpu params

Google-ML-Automation committed 209 days ago

5608cb7c

[pallas:sc] Skip `pl.multiple_of` test on old libtpu

superbobry committed 209 days ago

fff4d58d

[pallas:sc] Renamed the `pl.debug_check` test to match other SC tests

superbobry committed 209 days ago

8ade11bf

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

Google-ML-Automation committed 209 days ago

174025a6

Reverts 4ee9cb2bdf89fdd6e05b43b77e8cb593464e0a56

Google-ML-Automation committed 209 days ago

fc0c1d08

Switch ShapeDtypeStruct to use concrete_mesh when it's created with PartitionSpec during construction time.

yashk2810 committed 210 days ago

20e84d0a

make a "hello world" of pallas_call + pinned buffers

mattjj committed 210 days ago

50e2a612

Merge pull request #32586 from mattjj:hijax-compiled

Google-ML-Automation committed 210 days ago

a9b9cf3c

Add a fallback to `get_hlo_text()` to go via old code if AttributeError is hit.

Google-ML-Automation committed 210 days ago

1aa217b8

[buildcop] Fix acc_smem in matmul1

jakeharmon8 committed 210 days ago

5f351a0b

Fix typo in _one_hot

justinjfu committed 210 days ago

41bb2d7e

Older