From 6fc7f9855661eb942a136ceb7c27975d10d7e493 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 21:39:01 +0000 Subject: [PATCH 01/42] Remove nvidia_wheel_versions --- jaxlib/jax.bzl | 2 -- 1 file changed, 2 deletions(-) diff --git a/jaxlib/jax.bzl b/jaxlib/jax.bzl index 924400c6fe61..cb6d497c4618 100644 --- a/jaxlib/jax.bzl +++ b/jaxlib/jax.bzl @@ -22,7 +22,6 @@ load("@jax_wheel//:wheel.bzl", "WHEEL_VERSION") load("@jax_wheel_version_suffix//:wheel_version_suffix.bzl", "WHEEL_VERSION_SUFFIX") load("@local_config_cuda//cuda:build_defs.bzl", _cuda_library = "cuda_library", _if_cuda_is_configured = "if_cuda_is_configured") load("@local_config_rocm//rocm:build_defs.bzl", _if_rocm_is_configured = "if_rocm_is_configured", _rocm_library = "rocm_library") -load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS") load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION", "HERMETIC_PYTHON_VERSION_KIND") load("@rocm_external_test_deps//:external_deps.bzl", "EXTERNAL_DEPS") load("@rocm_prebuilt_test_deps//:external_deps.bzl", PREBUILT_EXTERNAL_DEPS = "EXTERNAL_DEPS") @@ -485,7 +484,6 @@ def _jax_wheel_impl(ctx): if ctx.attr.platform_version == "": fail("platform_version must be set to a valid cuda version for cuda wheels") args.add("--platform_version", ctx.attr.platform_version) # required for gpu wheels - args.add("--nvidia_wheel_versions_data", NVIDIA_WHEEL_VERSIONS) # required for gpu wheels if ctx.attr.enable_rocm: args.add("--enable-rocm", "True") if ctx.attr.platform_version == "": From 4c376bc1bd15fb6e76086af2cfddeb6a478447f1 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 21:42:44 +0000 Subject: [PATCH 02/42] Make jaxlib targets visible --- jaxlib/BUILD | 2 +- jaxlib/rocm/BUILD | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/jaxlib/BUILD b/jaxlib/BUILD index a434e7ebc119..134f364ffde6 100644 --- a/jaxlib/BUILD +++ b/jaxlib/BUILD @@ -41,7 +41,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//jax:internal"], + default_visibility = ["//visibility:public"], ) package_group( diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index 65a844a2e99f..50f8dada38f5 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -28,7 +28,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//:__subpackages__"], + default_visibility = ["//visibility:public"], ) cc_library( From 4bf1d24bac5c6dede3a085367259534c67e54ae2 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 22:12:01 +0000 Subject: [PATCH 03/42] hipblas typedef fix --- jaxlib/gpu/solver_interface.cc | 12 ++++++------ jaxlib/gpu/vendor.h | 12 ++++-------- jaxlib/rocm/BUILD | 4 ++++ 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/jaxlib/gpu/solver_interface.cc b/jaxlib/gpu/solver_interface.cc index 7ebcf7483652..e530e5bc1f2f 100644 --- a/jaxlib/gpu/solver_interface.cc +++ b/jaxlib/gpu/solver_interface.cc @@ -63,8 +63,8 @@ JAX_GPU_DEFINE_GETRF(gpuDoubleComplex, gpusolverDnZgetrf); JAX_GPU_DEFINE_GETRF_BATCHED(float, gpublasSgetrfBatched); JAX_GPU_DEFINE_GETRF_BATCHED(double, gpublasDgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasComplex, gpublasCgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasDoubleComplex, gpublasZgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuComplex, gpublasCgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuDoubleComplex, gpublasZgetrfBatched); #undef JAX_GPU_DEFINE_GETRF_BATCHED // QR decomposition: geqrf @@ -102,8 +102,8 @@ JAX_GPU_DEFINE_GEQRF(gpuDoubleComplex, gpusolverDnZgeqrf); JAX_GPU_DEFINE_GEQRF_BATCHED(float, gpublasSgeqrfBatched); JAX_GPU_DEFINE_GEQRF_BATCHED(double, gpublasDgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasComplex, gpublasCgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasDoubleComplex, gpublasZgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuComplex, gpublasCgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuDoubleComplex, gpublasZgeqrfBatched); #undef JAX_GPU_DEFINE_GEQRF_BATCHED // Householder transformations: orgqr @@ -302,8 +302,8 @@ JAX_GPU_DEFINE_SYEVD(gpuDoubleComplex, gpusolverDnZheevd); JAX_GPU_DEFINE_SYRK(float, gpublasSsyrk); JAX_GPU_DEFINE_SYRK(double, gpublasDsyrk); -JAX_GPU_DEFINE_SYRK(gpublasComplex, gpublasCsyrk); -JAX_GPU_DEFINE_SYRK(gpublasDoubleComplex, gpublasZsyrk); +JAX_GPU_DEFINE_SYRK(gpuComplex, gpublasCsyrk); +JAX_GPU_DEFINE_SYRK(gpuDoubleComplex, gpublasZsyrk); #undef JAX_GPU_DEFINE_SYRK // Singular Value Decomposition: gesvd diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 9adb1f11776f..4818a70e6941 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -457,6 +457,8 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; #elif defined(JAX_GPU_HIP) +#define HIPBLAS_V2 1 + // IWYU pragma: begin_exports #include "rocm/include/hip/hip_cooperative_groups.h" #include "rocm/include/hip/hip_runtime_api.h" @@ -477,17 +479,11 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; // MIOpen lib. Remove when MIOpen support is complete. #define MIOPEN_STATUS_SUCCESS 0 -typedef hipFloatComplex gpuComplex; +typedef hipComplex gpuComplex; typedef hipDoubleComplex gpuDoubleComplex; -#if TF_ROCM_VERSION >= 70000 -typedef hipFloatComplex gpublasComplex; +typedef hipComplex gpublasComplex; typedef hipDoubleComplex gpublasDoubleComplex; -#else -typedef hipblasComplex gpublasComplex; -typedef hipblasDoubleComplex gpublasDoubleComplex; -#endif // TF_ROCM_VERSION >= 70000 - typedef struct hipsolverHandle_* gpusolverDnHandle_t; typedef hipblasFillMode_t gpublasFillMode_t; typedef hipsolverFillMode_t gpusolverFillMode_t; diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index 50f8dada38f5..bae333592d16 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -425,6 +425,10 @@ rocm_nanobind_extension( "@xla//xla/ffi/api:ffi", "@xla//xla/python:safe_static_init", ], + linkopts = [ + "-L/opt/rocm/lib", + "-lamdhip64", + ], ) cc_library( From 47e08da48e471a327a26ceb0d059b8f4c68f700e Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Thu, 13 Nov 2025 21:50:02 +0000 Subject: [PATCH 04/42] No GPU fail --- jaxlib/rocm/BUILD | 1 + jaxlib/rocm/rocm_plugin_extension.cc | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index bae333592d16..3c0c93844487 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -558,6 +558,7 @@ rocm_nanobind_extension( enable_stub_generation = False, module_name = "rocm_plugin_extension", deps = [ + ":hip_gpu_kernel_helpers", ":py_client_gpu", "//jaxlib:kernel_nanobind_helpers", "//jaxlib/gpu:gpu_plugin_extension", diff --git a/jaxlib/rocm/rocm_plugin_extension.cc b/jaxlib/rocm/rocm_plugin_extension.cc index 05be1e81c858..e28e4927b81f 100644 --- a/jaxlib/rocm/rocm_plugin_extension.cc +++ b/jaxlib/rocm/rocm_plugin_extension.cc @@ -23,6 +23,7 @@ limitations under the License. #include "jaxlib/gpu/gpu_plugin_extension.h" #include "jaxlib/gpu/py_client_gpu.h" #include "jaxlib/kernel_nanobind_helpers.h" +#include "jaxlib/gpu/gpu_kernel_helpers.h" namespace nb = nanobind; @@ -96,6 +97,13 @@ nb::dict FfiHandlers() { return dict; } +int ROCmDeviceCount() { + int device_count = -1; + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipInit(0))); + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipGetDeviceCount(&device_count))); + return device_count; +} + } // namespace NB_MODULE(rocm_plugin_extension, m) { @@ -122,5 +130,6 @@ NB_MODULE(rocm_plugin_extension, m) { return device_ordinal; }, nb::arg("data_value")); + m.def("get_device_count", &ROCmDeviceCount); } } // namespace jax From 9180b93cbaf1644db261841a5b11ce45fdb3ac26 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 12 Feb 2026 14:06:16 -0800 Subject: [PATCH 05/42] Wrap HIP inline functions in anonymous namespaces in vendor.h --- jaxlib/gpu/vendor.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 4818a70e6941..31d554cf35b4 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -541,6 +541,7 @@ inline hipblasStatus_t gpublasCreate(gpublasHandle_t* handle) { return hipblasCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpublasCreate ::jax::hip::gpublasCreate #define gpublasSetStream hipblasSetStream #define gpublasSgeqrfBatched hipblasSgeqrfBatched @@ -597,6 +598,7 @@ inline hipsolverStatus_t gpusolverDnCreate(gpusolverDnHandle_t* handle) { return hipsolverCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpusolverDnCreate ::jax::hip::gpusolverDnCreate #define gpusolverDnSetStream hipsolverSetStream #define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo From aba01e099fafe09256d15b64737bd238c8830984 Mon Sep 17 00:00:00 2001 From: Dragoslav Sicarov Date: Tue, 10 Jun 2025 04:28:40 +0000 Subject: [PATCH 06/42] SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError --- jaxlib/gpu/vendor.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 31d554cf35b4..5c9098494886 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -794,7 +794,7 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPU_STREAM_NON_BLOCKING hipStreamNonBlocking #define gpuMalloc hipMalloc -#define gpuGetLastError hipGetLastError +#define gpuGetLastError hipExtGetLastError #define gpuGetErrorString hipGetErrorString #define gpuMemcpyAsync hipMemcpyAsync #define gpuMemcpyDeviceToDevice hipMemcpyDeviceToDevice From b1b7a8b4932c1a5cbc2a13b1120994c3929578e0 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Fri, 14 Nov 2025 19:43:09 +0000 Subject: [PATCH 07/42] Add shared utility function get_rocm_version to test_util.py --- jax/_src/test_util.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py index d688221e7881..b5a18aa6497d 100644 --- a/jax/_src/test_util.py +++ b/jax/_src/test_util.py @@ -25,6 +25,7 @@ import math import os +from pathlib import Path import platform import re import sys @@ -389,6 +390,15 @@ def supported_dtypes() -> set[DTypeLike]: def is_device_rocm() -> bool: return 'rocm' in xla_bridge.get_backend().platform_version +def get_rocm_version(): + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + version_path = Path(rocm_path) / ".info" / "version" + if not version_path.exists(): + raise FileNotFoundError(f"Expected ROCm version file at {version_path}") + version_str = version_path.read_text().strip() + major, minor, *_ = version_str.split(".") + return int(major), int(minor) + def is_device_cuda() -> bool: return 'cuda' in xla_bridge.get_backend().platform_version From bf79d6250965bf0511329329b9fdcff9921d5cd6 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Mon, 17 Nov 2025 19:34:29 +0000 Subject: [PATCH 08/42] Fix hipSparse CSR algorithm mappings for ROCm 7 --- jaxlib/gpu/vendor.h | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 5c9098494886..fb9c1f407145 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -778,10 +778,10 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPUSPARSE_INDEX_32I HIPSPARSE_INDEX_32I #define GPUSPARSE_INDEX_64I HIPSPARSE_INDEX_64I #define GPUSPARSE_DENSETOSPARSE_ALG_DEFAULT HIPSPARSE_DENSETOSPARSE_ALG_DEFAULT -#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_ALG_DEFAULT -#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_ALG_DEFAULT +#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_COOMV_ALG +#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_CSRMV_ALG1 +#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_COO_ALG1 +#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_CSR_ALG1 #define GPUSPARSE_INDEX_BASE_ZERO HIPSPARSE_INDEX_BASE_ZERO #define GPUSPARSE_OPERATION_NON_TRANSPOSE HIPSPARSE_OPERATION_NON_TRANSPOSE #define GPUSPARSE_OPERATION_TRANSPOSE HIPSPARSE_OPERATION_TRANSPOSE From c525fe7217a8f4c78ea09207d41541877b23dfce Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Thu, 20 Nov 2025 01:30:24 +0200 Subject: [PATCH 09/42] =?UTF-8?q?Fix=20v=5Fpages=20quantization=20and=20ad?= =?UTF-8?q?just=20test=20params=20for=20ROCm=20compatibilit=E2=80=A6=20(#5?= =?UTF-8?q?60)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/pallas/gpu_paged_attention_test.py | 71 +++++++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/tests/pallas/gpu_paged_attention_test.py b/tests/pallas/gpu_paged_attention_test.py index 1b778c787a6d..6a1d8de22a78 100644 --- a/tests/pallas/gpu_paged_attention_test.py +++ b/tests/pallas/gpu_paged_attention_test.py @@ -112,6 +112,64 @@ class PagedAttentionKernelTest(PallasBaseTest): def setUp(self): super().setUp() + def _estimate_shared_memory_bytes(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Estimate shared memory usage for paged attention kernel.""" + dtype_size = jnp.dtype(dtype).itemsize + # Approximate calculation based on kernel's memory usage + # Q block: block_h * head_dim + # K/V blocks: pages_per_compute_block * page_size * head_dim + # Plus accumulators and intermediate values + block_k = pages_per_compute_block * page_size + estimated = dtype_size * ( + block_h * head_dim + # Q + 2 * block_k * head_dim + # K and V + block_h * block_k + # logits/attention weights + block_h * 8 # accumulators (m, l, etc.) in float32 + ) + return estimated + + def _adjust_params_for_shared_memory(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Adjust parameters to fit within device shared memory limits. + + Uses XLA's DeviceDescription.shared_memory_per_block_optin() to query + the actual device capability rather than hardcoding values. + """ + try: + device = jax.local_devices()[0] + # Query XLA DeviceDescription for max shared memory per block + # This is exposed from stream_executor::DeviceDescription::shared_memory_per_block_optin() + max_smem = device.shared_memory_per_block_optin + except (AttributeError, IndexError): + # Fallback if XLA doesn't expose shared_memory_per_block_optin (older versions) + # or if no devices are available. Use conservative 48KB (safe for most GPUs). + max_smem = 48 * 1024 + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + # If within limits, no adjustment needed + if estimated <= max_smem: + return block_h, pages_per_compute_block, page_size + + # Try to reduce parameters to fit + while estimated > max_smem: + if pages_per_compute_block > 2: + pages_per_compute_block = pages_per_compute_block // 2 + elif page_size > 8: + page_size = page_size // 2 + elif block_h > 8: + block_h = block_h // 2 + else: + # Can't reduce further, will need to skip + return None, None, None + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + return block_h, pages_per_compute_block, page_size + @jtu.sample_product( dtype=(jnp.float16,), page_size=(8, 16, 32), @@ -201,6 +259,17 @@ def test_quantized_paged_attention( if (quant_dtype == jnp.float8_e4m3fn and not jtu.is_cuda_compute_capability_at_least("8.9")): self.skipTest("Skipping since float8_e4m3fn is not supported on < sm89") + + # Check and adjust parameters if needed to fit device limits for ROCm + if jtu.is_device_rocm(): + adjusted = self._adjust_params_for_shared_memory( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + if adjusted == (None, None, None): + self.skipTest("Cannot adjust parameters to fit ROCm device shared memory limits") + + block_h, pages_per_compute_block, page_size = adjusted + max_kv_len = 2048 seq_lens = np.asarray([3, 256, 513, 1023, 2048], dtype=jnp.int32) q, k_pages, v_pages, block_tables = _generate_qkv( @@ -218,7 +287,7 @@ def test_quantized_paged_attention( k_, k_scales = (_quantize(k_pages, quant_dtype) if quantize_k else (k_pages, None)) - v_, v_scales = (_quantize(k_pages, quant_dtype) + v_, v_scales = (_quantize(v_pages, quant_dtype) if quantize_v else (v_pages, None)) o = paged_attention.paged_attention( From f270959a5cebe095e9f120a3f03ef28f1e2fe3ef Mon Sep 17 00:00:00 2001 From: Aleksei <208770786+Arech8@users.noreply.github.com> Date: Wed, 26 Nov 2025 17:19:56 +0200 Subject: [PATCH 10/42] Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563) When jaxlib was built in debug more, an assertion in LLVM code that lazy-loads VHLO dialect could fire, since the code path could execute in a multi-threaded environment, and LLVM dialect repositories aren't thread safe to modify. This patch applies the same changes that upstream makes to fix this: https://github.com/jax-ml/jax/commit/48c876227a67840d7bcd44d80d16edbc0e910335 (this includes disabling a call to `jax_mlir_ext.enter_multi_threaded_execution(context)` in `mlir.py`. Presumably, the whole functionality related to `enter_multi_threaded_execution()` multithreaded checks isn't ready yet, and it was prematurely rolled into the production code. Manual testing --- .gitignore | 4 ++++ jax/_src/interpreters/mlir.py | 5 ++--- jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc | 1 - 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index d30c019b31e8..e55a9a5de101 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,7 @@ jax.iml /include/ /lib/ /share/ + +/compile_commands.json +/strace.txt +/external diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py index fe3ab7f910f0..5ea7f584e3f2 100644 --- a/jax/_src/interpreters/mlir.py +++ b/jax/_src/interpreters/mlir.py @@ -640,9 +640,8 @@ def make_ir_context() -> ir.Context: # multi threaded execution aborts the process if we try to register a new # dialect after this point. The dialect registry in a context is not thread # safe, and a fatal error is much better than a data race. - # jax_mlir_ext.enter_multi_threaded_execution(context) - # TODO(phawkins): clean up users who add their own dialects to JAX's contexts - # and enable this. + # if jaxlib_version >= (0, 8): + # jax_mlir_ext.enter_multi_threaded_execution(context) return context diff --git a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc index fc04f9ce1718..d6515d146e00 100644 --- a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc +++ b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc @@ -208,7 +208,6 @@ NB_MODULE(_jax_mlir_ext, m) { unwrap(registry)->insert(); unwrap(registry)->insert(); unwrap(registry)->insert(); - // For Mosaic GPU REGISTER_DIALECT(cf); REGISTER_DIALECT(gpu); From ced0e4d24619c4bcaf635a018d9c6d82d6c83f35 Mon Sep 17 00:00:00 2001 From: Aleksei <208770786+Arech8@users.noreply.github.com> Date: Wed, 26 Nov 2025 18:35:19 +0200 Subject: [PATCH 11/42] Add skip of test_is_finite() on Cuda (#565) (forgot this skip in the previous PR) --- tests/pallas/ops_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index 7b2b670a0db1..f4ee9013e2b9 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -1047,6 +1047,10 @@ def test_is_finite(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.test_device_matches(["cuda"]): + self.skipTest("Not tested on CUDA") # set this b/c this how the test was + # originally configured. Have no way to test cuda. + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( From 19c9a3ede0a41f5a026e46b01c003663b2259531 Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Mon, 15 Dec 2025 08:00:42 -0800 Subject: [PATCH 12/42] Add rocm test requirements file (#570) --- build/rocm-test-requirements.txt | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 build/rocm-test-requirements.txt diff --git a/build/rocm-test-requirements.txt b/build/rocm-test-requirements.txt new file mode 100644 index 000000000000..399237175957 --- /dev/null +++ b/build/rocm-test-requirements.txt @@ -0,0 +1,24 @@ +absl-py +build +cloudpickle +colorama>=0.4.4 +filelock +flatbuffers +hypothesis +mpmath>=1.3 +pillow>=10.4.0 +# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t +portpicker; python_version<"3.13" +pytest-xdist +pytest-json-report +pytest-html +pytest-csv +pytest-rerunfailures +pytest-html-merger +pytest-reportlog +wheel +rich +setuptools +matplotlib +opt-einsum +auditwheel From dba3b37d4e99cfbcd2727f3346376f42035f5f73 Mon Sep 17 00:00:00 2001 From: charleshofer Date: Mon, 15 Dec 2025 11:03:51 -0600 Subject: [PATCH 13/42] Let the unit tests use build.py for setting up Bazel commands for unit tests (#582) --- build/build.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/build/build.py b/build/build.py index 32a7b0016639..994e2df222ae 100755 --- a/build/build.py +++ b/build/build.py @@ -638,6 +638,10 @@ async def main(): ) if "rocm" in args.wheels: + if not args.configure_only: + print("ERROR: This repo is not used for building the ROCm JAX plugins. Please use the new plugin repo: https://github.com/ROCm/rocm-jax") + exit(1) + wheel_build_command_base.append("--config=rocm_base") wheel_build_command_base.append("--config=rocm") if clang_local: From 6c52a2721e9fd9037ca63adfbecb42e9a16610d7 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 13 Jan 2026 10:29:56 -0600 Subject: [PATCH 14/42] adding abort logic to rocm/jax (#590) --- conftest.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/conftest.py b/conftest.py index 1e6e17ab966a..13bb89c9ca0a 100644 --- a/conftest.py +++ b/conftest.py @@ -15,7 +15,10 @@ import os import pytest - +import json +import threading +import shutil +from datetime import datetime @pytest.fixture(autouse=True) def add_imports(doctest_namespace): From 2c1ba3a2ed4678fcee878f8bbb9598a48d0e4fff Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Wed, 14 Jan 2026 19:57:26 +0200 Subject: [PATCH 15/42] Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (#597) --- tests/pallas/ops_test.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index f4ee9013e2b9..81f7262e639e 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -1051,6 +1051,9 @@ def test_is_finite(self, dtype): self.skipTest("Not tested on CUDA") # set this b/c this how the test was # originally configured. Have no way to test cuda. + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( @@ -1103,6 +1106,9 @@ def test_is_finite_scalar(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( From 40323a1a54e2ea9b9fd731a61d41c8f7c76b14a8 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Wed, 14 Jan 2026 19:57:53 +0200 Subject: [PATCH 16/42] Fix shared memory limit check for ROCm in test_dot (#596) --- tests/pallas/ops_test.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index 81f7262e639e..1031f72c80b5 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -2170,12 +2170,20 @@ def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y): self.skipTest("bfloat16 type are not supported on GPU") # Check shared memory limit: Triton loads lhs + rhs into shared memory if jtu.is_device_rocm(): + # ROCm: use correct formula with dynamic limit from rocminfo dtype_size = jnp.dtype(dtype).itemsize - if (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size > get_rocm_shared_memory_limit(): + shared_mem_bytes = (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size + shared_mem_limit = get_rocm_shared_memory_limit() + if shared_mem_bytes > shared_mem_limit: self.skipTest("Shared memory size limit exceeded") - elif math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) > (256 * 256) * 2: - self.skipTest("Shared memory size limit exceeded") - if (jax.local_devices()[0].shared_memory_per_block_optin == 99 * 1024 and + else: + # NVIDIA: keep original check + if ( + math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) + > (256 * 256) * 2 + ): + self.skipTest("Shared memory size limit exceeded") + if (jax.local_devices()[0].device_kind == "NVIDIA L4" and dtype == jnp.float32 and lhs_and_rhs_shape in [ ((128, 16), (128, 256)), From cda6fb974eb277c0c92e5f78083d86f43dbdc0b0 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Wed, 14 Jan 2026 12:14:02 -0600 Subject: [PATCH 17/42] Fix Numpy signatures test (#598) Co-authored-by: Daniel Suo Co-authored-by: Jake VanderPlas --- tests/lax_numpy_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index 54911ab05be4..b04f6e87d94c 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6366,6 +6366,10 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], +<<<<<<< HEAD +======= + 'var': ['mean'], +>>>>>>> a3f8af53f (Fix Numpy signatures test (#598)) 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 55bbbdfb8d436b08c41de934af738dfff549452a Mon Sep 17 00:00:00 2001 From: Ruturaj4 Date: Sun, 18 Jan 2026 10:18:54 -0600 Subject: [PATCH 18/42] fix merge arts --- tests/lax_numpy_test.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index b04f6e87d94c..63624e1a4eef 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6366,10 +6366,7 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], -<<<<<<< HEAD -======= 'var': ['mean'], ->>>>>>> a3f8af53f (Fix Numpy signatures test (#598)) 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 4ec5da21fd36b793e672f7ee65cdc03f89b3cd92 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Thu, 22 Jan 2026 15:44:41 -0600 Subject: [PATCH 19/42] Enable RngShardingTests (#644) --- tests/array_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/array_test.py b/tests/array_test.py index 5e46ba709fee..fd67a31e28e5 100644 --- a/tests/array_test.py +++ b/tests/array_test.py @@ -1541,7 +1541,7 @@ class RngShardingTest(jtu.JaxTestCase): # tests that the PRNGs are automatically sharded as expected @parameterized.named_parameters(("3", 3), ("4", 4), ("5", 5)) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_1d(self, num_devices): @jax.jit def f(x): @@ -1575,7 +1575,7 @@ def f(x): "mesh_shape": mesh_shape, "pspec": pspec} for mesh_shape in [(3, 2), (4, 2), (2, 3)] for pspec in [P('x', None), P(None, 'y'), P('x', 'y')]) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_2d(self, mesh_shape, pspec): @jax.jit def f(x): From 6121864df7839743676606d8cd1fcfd4cb327da3 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 12 Feb 2026 15:18:40 -0800 Subject: [PATCH 20/42] Enable test_variadic_reduce_window on ROCm (#647) --- tests/lax_vmap_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py index 57fdfc4dda88..de260ec4ed93 100644 --- a/tests/lax_vmap_test.py +++ b/tests/lax_vmap_test.py @@ -759,6 +759,8 @@ def testSort(self, shape, dimension, arity, bdims, is_stable): # TODO Collapse # TODO Scatter + # b/183233858: variadic reduce-window not implemented on XLA:CUDA + @jtu.skip_on_devices("cuda") def test_variadic_reduce_window(self): # https://github.com/jax-ml/jax/discussions/9818 and # https://github.com/jax-ml/jax/issues/9837 From 8438836f1296a979499e8849dc43c5ec6d07b3eb Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 21/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index affcaf058c00..e7c2398c0c37 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -219,6 +220,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -592,6 +594,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -1050,6 +1053,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. From 3264ecbfc00c81536719d51044fd48892dd3d29d Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 22/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index e7c2398c0c37..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,11 +137,8 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") csr_matmul = sparse_csr._csr_matvec if len(bshape) == 1 else sparse_csr._csr_matmat tol = {np.float32: 2E-5, np.float64: 1E-12, np.complex64: 1E-5, @@ -220,11 +217,8 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") op = lambda M: M.T if transpose else M @@ -594,11 +588,8 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") tol = {np.float32: 2E-5, np.float64: 2E-14} @@ -1053,11 +1044,8 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") rng = sptu.rand_sparse(self.rng(), post=jnp.array) rng_b = jtu.rand_default(self.rng()) From 6a47a46764a24513c3185b047437fe4f605d0f5f Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 23/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index fa91f4f18185..6cf04543d74d 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1044,6 +1047,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From b4de9f9d7b43f2a57eac25190a89d20f47dbd853 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 24/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 6cf04543d74d..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1047,7 +1044,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 9c6283d3ad8d08484667ed9f5941905b2085b296 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 25/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index fa91f4f18185..6cf04543d74d 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1044,6 +1047,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 68f67b5818e2567ceae99d9aecbd02003e83da23 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 26/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 6cf04543d74d..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1047,7 +1044,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From cef5f06b5382ad786190fc0083460c1b46dfd689 Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Wed, 28 Jan 2026 14:00:42 -0800 Subject: [PATCH 27/42] Enable testMultivariateNormalSingularCovariance on ROCm (#666) --- tests/random_lax_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py index 0f317d81330a..ef0af2ec9ce3 100644 --- a/tests/random_lax_test.py +++ b/tests/random_lax_test.py @@ -968,7 +968,7 @@ def testMultivariateNormalCovariance(self): check_dtypes=False) @jtu.sample_product(method=['cholesky', 'eigh', 'svd']) - @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. + @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. ROCm supported def testMultivariateNormalSingularCovariance(self, method): # Singular covariance matrix https://github.com/jax-ml/jax/discussions/13293 mu = jnp.zeros((2,)) From 235f616cf35c2d8f8d65a3a3f37a07a60998394d Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Wed, 28 Jan 2026 16:38:17 -0600 Subject: [PATCH 28/42] Update Skip Reason Outputs (#663) --- jax/_src/test_util.py | 20 ++++++++++++++++---- tests/pallas/gpu_pallas_distributed_test.py | 13 ++++++++----- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py index b5a18aa6497d..0dee9c26e552 100644 --- a/jax/_src/test_util.py +++ b/jax/_src/test_util.py @@ -620,7 +620,14 @@ def skip_on_devices(*disabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = "Skipped on devices with tags: " + ", ".join(disabled_devices) + skip_messages = { + ("gpu",): "Skipped on all GPUs.", + ("cpu",): "Skipped on CPU.", + ("tpu",): "Skipped on TPU.", + ("cuda",): "Skipped on CUDA GPUs.", + ("rocm",): "Skipped on ROCm GPUs.", + } + skip_reason = skip_messages.get(disabled_devices) return _device_filter(lambda: not test_device_matches(disabled_devices), skip_reason) def run_on_devices(*enabled_devices, skip_reason=None): @@ -631,9 +638,14 @@ def run_on_devices(*enabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = ( - "Skipped unless running on devices with tags: " + ", ".join(enabled_devices) - ) + device_specific_skip_reasons = { + ("cpu",): "Skipped: CPU-only test.", + ("tpu",): "Skipped: TPU-only test.", + ("gpu",): "Skipped: GPU-only test.", + ("rocm",): "Skipped: ROCm-only test.", + ("cuda",): "Skipped: CUDA-only test.", + } + skip_reason = device_specific_skip_reasons.get(enabled_devices) return _device_filter(lambda: test_device_matches(enabled_devices), skip_reason) def device_supports_buffer_donation(): diff --git a/tests/pallas/gpu_pallas_distributed_test.py b/tests/pallas/gpu_pallas_distributed_test.py index 2f174b21ca23..7137e19d9ec0 100644 --- a/tests/pallas/gpu_pallas_distributed_test.py +++ b/tests/pallas/gpu_pallas_distributed_test.py @@ -126,14 +126,17 @@ def setUp(self): if jtu.test_device_matches(["rocm"]): self.skipTest("Mosaic GPU is not supported on ROCm.") + # Check mosaic support first (before GPU capability check) + if not mgpu.supports_cross_device_collectives(): + if jtu.test_device_matches(["rocm"]): + self.skipTest("Mosaic not supported on ROCm currently.") + else: + self.skipTest("NVSHMEM library unavailable.") if (not jtu.test_device_matches(["cuda"]) or not jtu.is_cuda_compute_capability_at_least("9.0")): self.skipTest("Only works on GPU with capability >= sm90") - if not mgpu.supports_cross_device_collectives(): - self.skipTest( - "Skip test since cross-device collectives are not supported" - " (either NVSHMEM is not available in multi-process mode, or mixed" - " mode is used).") + if jax.process_count() == 1: + self.skipTest("Test requires multiple processes.") if os.environ.get("XLA_PYTHON_CLIENT_ALLOCATOR", "") == "platform": self.skipTest("NVSHMEM doesn't work with the platform allocator.") From 9a5723d5b8c8c87012ac466369b8d9a52fed4f7b Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 29/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index fa91f4f18185..6cf04543d74d 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1044,6 +1047,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From ff5dfaad4bb1e872dbc0b1cbb9aa8b73b6304c43 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 30/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 6cf04543d74d..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1047,7 +1044,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 9bfe7d66374aa42be40098e9d2427c757f5a7f20 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Thu, 29 Jan 2026 13:12:07 -0600 Subject: [PATCH 31/42] Skip testCudaArrayInterfaceOnNonCudaFails on ROCm platform (#677) --- tests/array_interoperability_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py index d033a490f547..b14e9171f952 100644 --- a/tests/array_interoperability_test.py +++ b/tests/array_interoperability_test.py @@ -312,7 +312,7 @@ def testCudaArrayInterfaceOnNonCudaFails(self): self.assertFalse(hasattr(x, "__cuda_array_interface__")) with self.assertRaisesRegex( AttributeError, - "__cuda_array_interface__ is only defined for .*GPU buffers.", + "__cuda_array_interface__ is only defined for GPU buffers.", ): _ = x.__cuda_array_interface__ From ed839d83642a210fcb07a17dd350b877dc8b6e8b Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 32/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index fa91f4f18185..6cf04543d74d 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1044,6 +1047,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 8e1ebcbaf7b3570bb5139001f20de8d5462d04d9 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 33/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 6cf04543d74d..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1047,7 +1044,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 0609c7c0a4d3795a1a7ce5dea15fe09fceefa89e Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 34/42] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index fa91f4f18185..6cf04543d74d 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1044,6 +1047,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 7a2494edc65b6d0687e65f34478616cfbfdab226 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 35/42] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 6cf04543d74d..fa91f4f18185 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1047,7 +1044,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 486782bd2cb6b6e244181c9f89495a0eb0d3a622 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 6 Feb 2026 12:48:48 -0600 Subject: [PATCH 36/42] Remove 'mean' from unsupported params for jnp.var (#689) --- tests/lax_numpy_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index 63624e1a4eef..54911ab05be4 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6366,7 +6366,6 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], - 'var': ['mean'], 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 9ee94768ed5dbe99f0d3704cc4435c0e20f49d3d Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Mon, 9 Feb 2026 05:08:48 -0800 Subject: [PATCH 37/42] Skipping testEighTinyNorm due to hipSolver issues (#697) --- tests/linalg_test.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/linalg_test.py b/tests/linalg_test.py index 5776aec6cb10..56f2f91c5381 100644 --- a/tests/linalg_test.py +++ b/tests/linalg_test.py @@ -585,11 +585,12 @@ def testEighZeroDiagonal(self): np.linalg.norm(np.matmul(a, v) - w * v), 2.5 * eps * np.linalg.norm(a) ) + def testEighTinyNorm(self): - # Skip test on ROCm due to numerical error. Issue #34711 - # TODO(GulsumGudukbay): Unskip once fixed. if jtu.is_device_rocm(): - self.skipTest("Skipped on ROCm due to numerical error.") + # numerical errors seen as of ROCm 7.2 due to hipSolver issue + # TODO: re-enable the test once the hipSolver issue is fixed + self.skipTest("testEighNorm not supported on ROCm due to hipSOLVER issue") rng = jtu.rand_default(self.rng()) a = rng((300, 300), dtype=np.float32) eps = jnp.finfo(a.dtype).eps From f9528df20fe77325a27d0f645939deaee5da4e62 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Fri, 20 Feb 2026 11:17:09 -0600 Subject: [PATCH 38/42] Abort detection CI workflow (#688) --- .github/workflows/pytest_rocm_abort.yml | 162 ++++++++++++++++ .../wheel_tests_nightly_release_abort.yml | 53 ++++++ ci/run_pytest_rocm_abort.sh | 179 ++++++++++++++++++ conftest.py | 4 - 4 files changed, 394 insertions(+), 4 deletions(-) create mode 100644 .github/workflows/pytest_rocm_abort.yml create mode 100644 .github/workflows/wheel_tests_nightly_release_abort.yml create mode 100755 ci/run_pytest_rocm_abort.sh diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml new file mode 100644 index 000000000000..fcfdb06855d5 --- /dev/null +++ b/.github/workflows/pytest_rocm_abort.yml @@ -0,0 +1,162 @@ +# CI - Pytest ROCm (Abort Support) +# +# This workflow runs the ROCm tests with Pytest in ROCm GHCR containers, +# using the ROCm `pytest-abort` retry wrapper to detect/retry aborts/crashes. +# +# It can be triggered manually via workflow_dispatch or called by other workflows +# via workflow_call. +# +# It consists of the following job: +# run-tests: +# - Runs in ROCm container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest) +# - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from latest release. +# - Executes the `run_pytest_rocm_abort.sh` script, which installs wheel artifacts and +# runs the ROCm tests with Pytest under `pytest-abort-retry`. +name: CI - Pytest ROCm (Abort Support) + +on: + workflow_dispatch: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: choice + default: "linux-x86-64-4gpu-amd" + options: + - "linux-x86-64-1gpu-amd" + - "linux-x86-64-4gpu-amd" + - "linux-x86-64-8gpu-amd" + python: + description: "Which Python version to use?" + type: choice + default: "3.11" + options: + - "3.11" + - "3.12" + rocm-version: + description: "Which ROCm version to test?" + type: choice + default: "7.2.0" + options: + - "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: choice + default: "head" + options: + - "head" + - "pypi_latest" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + type: choice + default: '0' + options: + - '0' + - '1' + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + type: string + default: 'gs://jax-nightly-artifacts/latest' + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + workflow_call: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: string + default: "linux-x86-64-4gpu-amd" + python: + description: "Which Python version to use?" + type: string + default: "3.11" + rocm-version: + description: "Which ROCm version to test?" + type: string + default: "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: string + default: "head" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + default: '0' + type: string + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + default: 'gs://jax-nightly-artifacts/latest' + type: string + +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-tests: + defaults: + run: + # Set the shell to bash as GitHub actions run with /bin/sh by default + shell: bash + runs-on: ${{ inputs.runner }} + continue-on-error: true + # Run in ROCm GHCR container with GPU access + container: + image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings + name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') || + (contains(inputs.runner, '4gpu') && '4gpu') || + (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}" + + env: + JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" + JAXCI_PYTHON: "python${{ inputs.python }}" + JAXCI_ENABLE_X64: "0" + + steps: + - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0 + with: + persist-credentials: false + - name: Download JAX ROCm wheels + uses: ./.github/actions/download-jax-rocm-wheels + with: + python: ${{ inputs.python }} + rocm-version: ${{ inputs.rocm-version }} + jaxlib-version: ${{ inputs.jaxlib-version }} + skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} + gcs_download_uri: ${{ inputs.gcs_download_uri }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Install Python dependencies + run: | + $JAXCI_PYTHON -m pip install uv~=0.5.30 + $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt + # Halt for testing + - name: Wait For Connection + uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c + with: + halt-dispatch-input: ${{ inputs.halt-for-connection }} + - name: Run Pytest ROCm tests (abort support) + timeout-minutes: 180 + run: ./ci/run_pytest_rocm_abort.sh + - name: Upload pytest results to artifact + if: always() + uses: actions/upload-artifact@v4 + with: + name: logs_abort + path: | + logs_abort/ + if-no-files-found: warn + retention-days: 2 + overwrite: true diff --git a/.github/workflows/wheel_tests_nightly_release_abort.yml b/.github/workflows/wheel_tests_nightly_release_abort.yml new file mode 100644 index 000000000000..518306956eff --- /dev/null +++ b/.github/workflows/wheel_tests_nightly_release_abort.yml @@ -0,0 +1,53 @@ +# CI - Wheel Tests (Nightly/Release) (ROCm abort only) +# +# This workflow runs only the ROCm wheel tests using the abort/retry wrapper workflow. +name: CI - Wheel Tests (Nightly/Release) (ROCm abort only) + +on: + workflow_dispatch: + inputs: + gcs_download_uri: + description: "GCS location URI from where the artifacts should be downloaded" + required: true + default: 'gs://jax-nightly-artifacts/latest' + type: string + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)" + required: true + default: '0' + type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection? (yes/no)' + required: false + default: 'no' + type: string + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: true +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-pytest-rocm: + uses: ./.github/workflows/pytest_rocm_abort.yml + strategy: + fail-fast: false # don't cancel all jobs on failure + matrix: + runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] + python: ["3.11", "3.12", "3.13", "3.14"] + rocm: [ + {version: "7.2.0", tag: "rocm720"}, + ] + name: "Pytest ROCm abort (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" + with: + runner: ${{ matrix.runner }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.version }} + rocm-tag: ${{ matrix.rocm.tag }} + jaxlib-version: "head" + skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} + gcs_download_uri: ${{inputs.gcs_download_uri}} + halt-for-connection: ${{inputs.halt-for-connection}} diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh new file mode 100755 index 000000000000..f17e7184b3bf --- /dev/null +++ b/ci/run_pytest_rocm_abort.sh @@ -0,0 +1,179 @@ +#!/bin/bash +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Runs Pytest ROCm tests (with ROCm pytest-abort retry wrapper). +# Requires the jaxlib and ROCm plugin wheels to be present inside $JAXCI_OUTPUT_DIR (../dist) +# +# -e: abort script if one command fails +# -u: error if undefined variable used +# -x: log all commands +# -o history: record shell history +# -o allexport: export all functions and variables to be available to subscripts +set -exu -o history -o allexport + +source ci/envs/default.env + +# Install jaxlib and ROCm plugin wheels inside the $JAXCI_OUTPUT_DIR directory +echo "Installing wheels locally..." +source ./ci/utilities/install_wheels_locally.sh + +# Print all the installed packages +echo "Installed packages:" +"$JAXCI_PYTHON" -m uv pip freeze + +"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))" + +rocm-smi + +# ============================================================================== +# Set up the generic test environment variables +# ============================================================================== +export PY_COLORS=1 +export JAX_SKIP_SLOW_TESTS=true +export NCCL_DEBUG=WARN +export TF_CPP_MIN_LOG_LEVEL=0 +export JAX_ENABLE_X64="$JAXCI_ENABLE_X64" + +# ============================================================================== +# Calculate the optimal number of parallel processes for pytest +# This will be the minimum of: GPU capacity, CPU core count, and a system RAM limit. +# ============================================================================== + +export gpu_count=$(rocminfo | egrep -c "Device Type:\\s+GPU") +echo "Number of GPUs detected: $gpu_count" + +# Query GPU 0 memory using rocm-smi +export memory_per_gpu_mib=$(rocm-smi -d 0 --showmeminfo vram | grep -i "vram total" | awk '{print int($NF/1024/1024)}' | head -1) +echo "Reported memory per GPU: $memory_per_gpu_mib MiB" + +# Convert effective memory from MiB to GiB. +export memory_per_gpu_gib=$((memory_per_gpu_mib / 1024)) +echo "Effective memory per GPU: $memory_per_gpu_gib GiB" + +# Allow 2 GiB of GPU RAM per test. +export max_tests_per_gpu=$((memory_per_gpu_gib / 2)) +echo "Max tests per GPU (assuming 2GiB/test): $max_tests_per_gpu" + +export num_processes=$((gpu_count * max_tests_per_gpu)) +echo "Initial number of processes based on GPU capacity: $num_processes" + +export num_cpu_cores=$(nproc) +echo "Number of CPU cores available: $num_cpu_cores" + +# Reads total memory from /proc/meminfo (in KiB) and converts to GiB. +export total_ram_gib=$(awk '/MemTotal/ {printf \"%.0f\", $2/1048576}' /proc/meminfo) +echo "Total system RAM: $total_ram_gib GiB" + +# Set a safety limit for system RAM usage, e.g., 1/6th of total. +export host_memory_limit=$((total_ram_gib / 6)) +echo "Host memory process limit (1/6th of total RAM): $host_memory_limit" + +if [[ $num_cpu_cores -lt $num_processes ]]; then + num_processes=$num_cpu_cores + echo "Adjusting num_processes to match CPU core count: $num_processes" +fi + +if [[ $host_memory_limit -lt $num_processes ]]; then + num_processes=$host_memory_limit + echo "Adjusting num_processes to match host memory limit: $num_processes" +fi + +if [[ 16 -lt $num_processes ]]; then + num_processes=16 + echo "Reducing num_processes to $num_processes" +fi + +echo "Final number of processes to run: $num_processes" + +export JAX_ENABLE_ROCM_XDIST="$gpu_count" +export XLA_PYTHON_CLIENT_ALLOCATOR=platform +export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1 --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_enable_command_buffer=" + +# Disable core dumps just in case +ulimit -c 0 + +# Keep deselected tests in one place for the abort wrapper. +ROCM_PYTEST_DESELECT_ARGS=( + --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data + --deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices + --deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric +) + +# --max-runs: retry the entire pytest run up to N times on abort/crash. +# --max-worker-restart: restart crashed xdist workers up to N times. +# --maxfail: stop the run after N test failures. +rocm_test_cmd() { + local abort_flag="${1:-0}" + shift + if [[ "$abort_flag" == "1" ]]; then + pytest-abort-retry --max-runs 3 --clear-crash-log -- "$@" + else + "$@" + fi +} + +rocm_log_tail_on_failure() { + local logfile="$1" + local status="$2" + if [[ "$status" -ne 0 ]]; then + echo "Pytest failed (exit=$status). Showing last 200 lines of $logfile:" + tail -n 200 "$logfile" || true + else + echo "Pytest output saved to $logfile (uploaded as artifact)." + fi +} + +rocm_install_extra_requirements() { + if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then + cd "$GITHUB_WORKSPACE" + fi + + # Install extra requirements. + "$JAXCI_PYTHON" -m uv pip install pytest-timeout pytest-html pytest-csv pytest-json-report pytest-abort +} + +rocm_install_extra_requirements + +echo "Running ROCm tests (with abort/retry wrapper)..." +mkdir -p logs_abort +logfile="logs_abort/jax_ToT_UT_abort.log" + +# pytest-abort output directories (must be set before running pytest). +export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running" +export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl" +mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR" + +set +e +rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart=200 --tb=short --timeout=1200 --timeout-method=thread tests \ + "${ROCM_PYTEST_DESELECT_ARGS[@]}" \ + --json-report \ + --json-report-file=logs_abort/tests-report-abort.json \ + --csv=logs_abort/tests-report-abort.csv \ + --html=logs_abort/tests-report-abort.html \ + --self-contained-html \ + >"$logfile" 2>&1 +pytest_status=$? +set -e +rocm_log_tail_on_failure "$logfile" "$pytest_status" + +echo "Postprocessing reports with crashed tests..." +pytest-abort-postprocess \ + --crash-log "$PYTEST_ABORT_CRASHED_TESTS_LOG" \ + --json-report logs_abort/tests-report-abort.json \ + --html-report logs_abort/tests-report-abort.html \ + --csv-report logs_abort/tests-report-abort.csv \ + >>"$logfile" 2>&1 + +exit "$pytest_status" diff --git a/conftest.py b/conftest.py index 13bb89c9ca0a..c313ecaf31c1 100644 --- a/conftest.py +++ b/conftest.py @@ -15,10 +15,6 @@ import os import pytest -import json -import threading -import shutil -from datetime import datetime @pytest.fixture(autouse=True) def add_imports(doctest_namespace): From e97e0ed313d8a0f27bedfd7e842537cc25ebba9d Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 24 Feb 2026 09:59:51 -0600 Subject: [PATCH 39/42] Abort-Detection: Fix halt-for-connection input (#712) --- .github/workflows/pytest_rocm_abort.yml | 13 +++++++++++++ ci/run_pytest_rocm_abort.sh | 7 +++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml index fcfdb06855d5..854b6681978b 100644 --- a/.github/workflows/pytest_rocm_abort.yml +++ b/.github/workflows/pytest_rocm_abort.yml @@ -64,6 +64,10 @@ on: description: 'Should this workflow run wait for a remote connection?' type: string default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' workflow_call: inputs: runner: @@ -94,6 +98,14 @@ on: description: "GCS location prefix from where the artifacts should be downloaded" default: 'gs://jax-nightly-artifacts/latest' type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' permissions: {} @@ -123,6 +135,7 @@ jobs: JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" JAXCI_PYTHON: "python${{ inputs.python }}" JAXCI_ENABLE_X64: "0" + MAX_WORKER_RESTART: "${{ inputs['max-worker-restart'] }}" steps: - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0 diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh index f17e7184b3bf..6c6ab6d18281 100755 --- a/ci/run_pytest_rocm_abort.sh +++ b/ci/run_pytest_rocm_abort.sh @@ -73,7 +73,7 @@ export num_cpu_cores=$(nproc) echo "Number of CPU cores available: $num_cpu_cores" # Reads total memory from /proc/meminfo (in KiB) and converts to GiB. -export total_ram_gib=$(awk '/MemTotal/ {printf \"%.0f\", $2/1048576}' /proc/meminfo) +export total_ram_gib=$(awk '/MemTotal/ {printf "%.0f", $2/1048576}' /proc/meminfo) echo "Total system RAM: $total_ram_gib GiB" # Set a safety limit for system RAM usage, e.g., 1/6th of total. @@ -150,13 +150,16 @@ echo "Running ROCm tests (with abort/retry wrapper)..." mkdir -p logs_abort logfile="logs_abort/jax_ToT_UT_abort.log" +# Allow the workflow to override worker restart limit. +max_worker_restart="${MAX_WORKER_RESTART:-50}" + # pytest-abort output directories (must be set before running pytest). export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running" export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl" mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR" set +e -rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart=200 --tb=short --timeout=1200 --timeout-method=thread tests \ +rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart="$max_worker_restart" --tb=short --timeout=1200 --timeout-method=thread tests \ "${ROCM_PYTEST_DESELECT_ARGS[@]}" \ --json-report \ --json-report-file=logs_abort/tests-report-abort.json \ From db0dd97b8c1bc29a8613b52a9682dd4055dd98c2 Mon Sep 17 00:00:00 2001 From: Yanyao Wang Date: Tue, 17 Mar 2026 17:40:53 -0500 Subject: [PATCH 40/42] fix: add rocm_sysdeps/lib to wheel RUNPATH (#737) Signed-off-by: Yanyao Wang --- jaxlib/rocm/rocm_rpath.bzl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl index 354082c0d200..1da2d640ca42 100644 --- a/jaxlib/rocm/rocm_rpath.bzl +++ b/jaxlib/rocm/rocm_rpath.bzl @@ -35,8 +35,11 @@ _ROCM_LINK_ONLY = "@local_config_rocm//rocm:link_only" _WHEEL_RPATHS = [ "-Wl,-rpath,$$ORIGIN/../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,$$ORIGIN/../../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,/opt/rocm/lib", + "-Wl,-rpath,/opt/rocm/lib/rocm_sysdeps/lib", ] def _wheel_features(): From ddcf4f412bbbd7fd8b6b8832da0bac4a95e7eba0 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Wed, 8 Apr 2026 08:58:12 -0500 Subject: [PATCH 41/42] Temporarily disable the cron trigger for the cont. wheel tests workflow --- .github/workflows/wheel_tests_continuous.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 657eccf4c248..4d0255e3b135 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -35,8 +35,8 @@ permissions: actions: read on: - schedule: - - cron: "0 */3 * * *" # Run once every 3 hours + # schedule: + # - cron: "0 */3 * * *" # Run once every 3 hours workflow_dispatch: # allows triggering the workflow run manually concurrency: From 37cbb3d6ff0b994fd318d6a892b5ba5391a7d41d Mon Sep 17 00:00:00 2001 From: Ken Date: Tue, 28 Apr 2026 17:51:32 -0500 Subject: [PATCH 42/42] Skip core_test reduction autotuner failures on Navi/RDNA GPUs (gfx11xx) XLA's GPU reduction autotuner has no valid tile configurations for gfx1151 (RDNA3.5/STRIX-HALO) because it uses 32-thread wavefronts (RDNA) vs the 64-thread wavefronts (CDNA) the autotuner was tuned for. JIT compilation triggers autotuning which fails with: NOT_FOUND: No valid config found! Affected tests: - CoreTest::test_simple_jit - CoreTest::test_tracer_reprs - CoreTest::test_verbose_tracer_reprs Follows the same skip pattern as ROCm/jax#720 (fused attention / TF32 skips). Detection: ROCm device + "Radeon" in device_kind. Tracking: ROCM-21429 Long-term fix: add valid XLA reduction autotuner tile configs for WarpSize==32 in xla/service/gpu/reduction_utils.cc. --- tests/core_test.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/core_test.py b/tests/core_test.py index e7e090f28d21..a99f09501037 100644 --- a/tests/core_test.py +++ b/tests/core_test.py @@ -256,6 +256,13 @@ def bar(y): assert jvp(foo, (1.0,), (2.0,)) == (1.0, 2.0) def test_simple_jit(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") def foo(x): if x.shape == (): return x + 1. @@ -403,6 +410,13 @@ def f(y): self.assertEqual(e1.outvars[0].aval.shape, (3, 3)) # only primal out shape def test_tracer_reprs(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") def f(x): nonlocal x_repr x_repr = repr(x) @@ -425,6 +439,13 @@ def f(x): self.assertEqual(x_repr, "JVPTracer(primal=[0. 1. 2.], tangent=VmapTracer(aval=float32[3], batched=float32[3,3]))") def test_verbose_tracer_reprs(self): + # XLA reduction autotuner produces no valid tile configs for gfx1151 (RDNA3.5, + # 32-thread wavefronts). JIT compilation triggers autotuning which fails with + # "NOT_FOUND: No valid config found!". See ROCM-21429. + if (jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + self.skipTest( + "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)") # Verbose reprs, avaiable via tracer._pretty_print() def f(x): nonlocal x_repr