ROCm/jax

Pull Requests Commits

Fix Numpy signatures test (#598)

magaonka-amd committed 3 days ago

c979baff

Fix shared memory limit check for ROCm in test_dot (#596)

phambinhfin committed 3 days ago

cbfb8422

Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (#597)

phambinhfin committed 3 days ago

becc59ee

adding abort logic to rocm/jax (#590)

gulsumgudukbay committed 3 days ago

5a2d899d

Let the unit tests use build.py for setting up Bazel commands for unit tests (#582)

charleshofer committed 3 days ago

11a084ba

Add rocm test requirements file (#570)

AratiGanesh committed 3 days ago

3a63ef39

Add skip of test_is_finite() on Cuda (#565)

Arech8 committed 3 days ago

d9103b3a

Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563)

Arech8 committed 3 days ago

04b3d823

Fix v_pages quantization and adjust test params for ROCm compatibilit… (#560)

phambinhfin committed 3 days ago

c3cf5d37

Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin committed 3 days ago

d6acb720

Add shared utility function get_rocm_version to test_util.py

charleshofer committed 3 days ago

c3e23572

SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd committed 3 days ago

ce432993

Wrap HIP inline functions in anonymous namespaces in vendor.h

mminutoli committed 3 days ago

a3338bb6

charleshofer committed 3 days ago

e8748279

hipblas typedef fix

charleshofer committed 3 days ago

d3f32d9b

Make jaxlib targets visible

charleshofer committed 3 days ago

6ba78d22

Remove nvidia_wheel_versions

charleshofer committed 3 days ago

2faed86f

Update rules_cc version to 0.2.9

akuegel committed 3 days ago

8b3bc0b4

[Mosaic:GPU] Use 2 GPUs for collective metadata test

PatriosTheGreat committed 3 days ago

1ff638f9

[Mosaic GPU] Add support for tcgen05 MMAs that are both sparse and block-scaled

apaszke committed 3 days ago

4553a01e

Re-enable CI for B200 MGPU test + increase shard_count

apaszke committed 3 days ago

cf870cbe

[pallas:mosaic] Renamed the `KernelType` enum to `CoreType`

superbobry committed 3 days ago

1d0504b1

[Pallas:MGPU] Expose async SMEM->TMEM copies for arbitrary data

apaszke committed 3 days ago

757c92db

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

Google-ML-Automation committed 3 days ago

4a28ce12

[Pallas] DCE all jaxprs before lowering

sharadmv committed 3 days ago

8fdc565d

Support atan2 in Pallas/Mosaic.

yarongmu-google committed 3 days ago

3d0ed0f8

[Rollback] remove jax_experimental_unsafe_xla_runtime_errors

mattjj committed 3 days ago

bb04e7a6

update_vma on AbstractRef updates its inner aval

Google-ML-Automation committed 3 days ago

11a3faf2

Pass overridden XLA repository location to CUDA non-RBE tests.

ybaturina committed 4 days ago

575ce86e

[Mosaic GPU] Remove legacy custom call API.

allanrenucci committed 4 days ago

5f99745e

Newer Older