ROCm/jax

Pull Requests Commits

- Fix `gpu_ops_test.py`: Implement block size reduction code. Dynamically reduce the block sizes until it fits the memory.

AratiGanesh committed 20 days ago

cff5d16a

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

Arech8 committed 53 days ago

Verified 861cc763

Implement Pallas lowering of lax.is_finite() for Triton (#564)

Arech8 committed 53 days ago

Verified c972c55d

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

Arech8 committed 53 days ago

Verified 339bf1e8

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

phambinhfin committed 60 days ago

Verified f893175a

Fix ROCm lib names copy pattern to use wildcard version

phambinhfin committed 60 days ago

709963b8

Make nvidia version data optional for ROCm builds

phambinhfin committed 60 days ago

b5762449

Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin committed 60 days ago

28e04303

Add shared utility function get_rocm_version to test_util.py

charleshofer committed 60 days ago

aac9e9d2

SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd committed 60 days ago

954fde6a

Wrap HIP inline functions in anonymous namespaces in vendor.h

AratiGanesh committed 60 days ago

13a4d047

charleshofer committed 60 days ago

f4df93f6

hipblas typedef fix

charleshofer committed 60 days ago

9ed3820b

Make jaxlib targets visible

charleshofer committed 60 days ago

3f0a5750

Remove nvidia_wheel_versions

charleshofer committed 60 days ago

f037da4e

Reverts f44bc6e37f770a9d79d24061f7960235a2892407

jakeharmon8 committed 95 days ago

403977d2

Prepare for JAX release 0.8.0

jakeharmon8 committed 95 days ago

b1ce4c57

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

justinjfu committed 95 days ago

07dad8e8

Update XLA hash for JAX release

jakeharmon8 committed 95 days ago

4f06e7a6

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

hawkinsp committed 95 days ago

1bd3402c

Clean up dead code.

brianwa84 committed 95 days ago

b942e894

Remove cusolver_potrf_ffi from the list of covered backwards compatibility targets.

gnecula committed 95 days ago

ac87cbe0

[Pallas:MGPU] Expose multicast TMA stores

apaszke committed 95 days ago

69fd2052

[Pallas:MGPU] Add an all-gather implementation

apaszke committed 95 days ago

8eed7dc5

Add tpu_sc.subcore_barrier primitive.

brianwa84 committed 95 days ago

6aae6f43

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

brianwa84 committed 95 days ago

8674d771

[Pallas:MGPU] Add ReducedLayout support to inline_mgpu params

Google-ML-Automation committed 95 days ago

5608cb7c

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

superbobry committed 95 days ago

fff4d58d

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

superbobry committed 95 days ago

8ade11bf

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

Google-ML-Automation committed 95 days ago

174025a6

Older