ROCm/jax

Pull Requests Commits

Add rocm test requirements file (#570)

AratiGanesh committed 8 days ago

3a63ef39

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

Arech8 committed 8 days ago

d9103b3a

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

Arech8 committed 8 days ago

04b3d823

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

phambinhfin committed 8 days ago

c3cf5d37

Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin committed 8 days ago

d6acb720

Add shared utility function get_rocm_version to test_util.py

charleshofer committed 8 days ago

c3e23572

SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd committed 8 days ago

ce432993

Wrap HIP inline functions in anonymous namespaces in vendor.h

mminutoli committed 8 days ago

a3338bb6

charleshofer committed 8 days ago

e8748279

hipblas typedef fix

charleshofer committed 8 days ago

d3f32d9b

Make jaxlib targets visible

charleshofer committed 8 days ago

6ba78d22

Remove nvidia_wheel_versions

charleshofer committed 8 days ago

2faed86f

Update rules_cc version to 0.2.9

akuegel committed 8 days ago

8b3bc0b4

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

PatriosTheGreat committed 9 days ago

1ff638f9

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

apaszke committed 9 days ago

4553a01e

Re-enable CI for B200 MGPU test + increase shard_count

apaszke committed 9 days ago

cf870cbe

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

superbobry committed 9 days ago

1d0504b1

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

apaszke committed 9 days ago

757c92db

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

Google-ML-Automation committed 9 days ago

4a28ce12

[Pallas] DCE all jaxprs before lowering

sharadmv committed 9 days ago

8fdc565d

Support atan2 in Pallas/Mosaic.

yarongmu-google committed 9 days ago

3d0ed0f8

[Rollback] remove jax_experimental_unsafe_xla_runtime_errors

mattjj committed 9 days ago

bb04e7a6

update_vma on AbstractRef updates its inner aval

Google-ML-Automation committed 9 days ago

11a3faf2

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

ybaturina committed 9 days ago

575ce86e

[Mosaic GPU] Remove legacy custom call API.

allanrenucci committed 9 days ago

5f99745e

[mosaic] Added `GetFuncWithCoreType`

superbobry committed 9 days ago

5c11726b

[Mosaic] Support erf.

WindQAQ committed 9 days ago

6e45cca4

Update rules_python version to 1.6.3 and gRPC to 1.78.0 in MODULE.bazel.

ybaturina committed 9 days ago

fe997775

Merge pull request #35236 from superbobry:main

Google-ML-Automation committed 9 days ago

1da5af44

[checkify] remove jax_experimental_unsafe_xla_runtime_errors

mattjj committed 9 days ago

bf610473

Newer Older