Enable testCaiToJax in array interoperability test on ROCm platform

This change enables the tests/array_interoperability_test.py::CudaArrayInterfaceTest::testCaiToJax on ROCm GPUs:

1. jaxlib/py_array.cc: Add ROCm platform check alongside CUDA for
   __cuda_array_interface__ property support.

2. jax/_src/numpy/array_constructors.py:
   - Add ROCm plugin extension discovery

3. tests/array_interoperability_test.py: Change testCaiToJax to use
   @jtu.run_on_devices("gpu") to run on both CUDA and ROCm.

Author: magaonka-amd

jax/_src/numpy/array_constructors.py

@@ -44,6 +44,24 @@
   else:
     break
 
+# Dynamically find and load ROCm plugin extension
+rocm_plugin_extension = None
+try:
+  from importlib.metadata import distributions
+  for dist in distributions():
+    name = dist.metadata.get('Name', '')
+    if name.startswith('jax-rocm') and name.endswith('-plugin'):
+      module_name = name.replace('-', '_')
+      try:
+        rocm_plugin_extension = importlib.import_module(
+            f'{module_name}.rocm_plugin_extension'
+        )
+        break
+      except ImportError:
+        continue
+except Exception:
+  pass
+
 
 def _supports_buffer_protocol(obj):
   try:
@@ -218,11 +236,17 @@ def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
       object = object.__jax_array__()
     elif hasattr(object, '__cuda_array_interface__'):
       cai = object.__cuda_array_interface__
-      backend = xla_bridge.get_backend("cuda")
-      if cuda_plugin_extension is None:
+      backend = xla_bridge.get_backend()
+      if 'rocm' in backend.platform_version.lower():
+        gpu_plugin_extension = rocm_plugin_extension
+      elif 'cuda' in backend.platform_version.lower():
+        gpu_plugin_extension = cuda_plugin_extension
+      else:
+        gpu_plugin_extension = None
+      if gpu_plugin_extension is None:
         device_id = None
       else:
-        device_id = cuda_plugin_extension.get_device_ordinal(cai["data"][0])
+        device_id = gpu_plugin_extension.get_device_ordinal(cai["data"][0])
       object = xc._xla.cuda_array_interface_to_buffer(
           cai=cai, gpu_backend=backend, device_id=device_id)
 

jaxlib/py_array.cc

@@ -985,9 +985,10 @@ nb::dict PyArray::CudaArrayInterface() {
   ifrt::Array* ifrt_array = arr.ifrt_array();
   std::optional<xla::Shape>& scratch = arr.GetStorage().dynamic_shape;
   auto* pjrt_buffer = GetPjrtBuffer(ifrt_array);
-  if (pjrt_buffer->client()->platform_id() != xla::CudaId()) {
+  if (pjrt_buffer->client()->platform_id() != xla::CudaId() &&
+      pjrt_buffer->client()->platform_id() != xla::RocmId()) {
     throw nb::attribute_error(
-        "__cuda_array_interface__ is only defined for NVidia GPU buffers.");
+        "__cuda_array_interface__ is only defined for GPU buffers.");
   }
   if (pjrt_buffer->IsTuple()) {
     throw nb::attribute_error(

tests/array_interoperability_test.py

@@ -392,7 +392,7 @@ def testCuPyToJax(self, shape, dtype):
     shape=all_shapes,
     dtype=jtu.dtypes.supported(cuda_array_interface_dtypes),
   )
-  @jtu.run_on_devices("cuda")
+  @jtu.run_on_devices("gpu")
   def testCaiToJax(self, shape, dtype):
     dtype = np.dtype(dtype)
 
@@ -401,7 +401,7 @@ def testCaiToJax(self, shape, dtype):
 
     # using device with highest device_id for testing the correctness
     # of detecting the device id from a pointer value
-    device = jax.devices('cuda')[-1]
+    device = jax.devices('gpu')[-1]
     with jax.default_device(device):
       y = jnp.array(x, dtype=dtype)
       # TODO(parkers): Remove after setting 'stream' properly below.