Repository: google/jax

Issue:
Support __cuda_array_interface__ on GPU 
https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html

It would not be hard to make `DeviceArray` implement this interface on GPU.

It would be slightly harder to support wrapping a `DeviceArray` around an existing CUDA array, but not that hard.


Discussion:
I think even just supporting one of the directions (i.e. making `DeviceArray` implement this interface on GPU) would already be a great addition.
I would be happy to help, but I am not sure where to find the pointer to GPU memory / what else to pay attention to.
TensorFlow now supports dlpack: https://github.com/VoVAllen/tf-dlpack/issues/3
PR #2133 added `__cuda_array_interface__` export. You'll need a jaxlib built from GitHub head or you'll need to wait for us to make another jaxlib wheel release.

Because of https://github.com/pytorch/pytorch/issues/32868 you can't directly import the resulting arrays to PyTorch. But because of https://github.com/cupy/cupy/issues/2616 you can "launder" the array via CuPy and into PyTorch if you want.

(Another option for interoperability is DLPack, which JAX supports at Github head, in both directions.)
Could this be reopened until import support is added as well?
I don't follow. We support both directions, I believe?

Edit: I am wrong, apparently we don't support imports.
> Edit: I am wrong, apparently we don't support imports.

Yeah this need came up again recently ( cc @leofang @quasiben ).
Although note that DLPack imports should work, so that's an option if the exporter supports DLPack.
Thanks John! Yeah we just finished a GPU Hackathon, and a few of our teams evaluating JAX asked us why JAX can't work with other libraries like CuPy and PyTorch _bidirectionally_. It'd be very useful, say, to do autograd in JAX, postprocess in CuPy, then bring it back to JAX.
Also: I haven't tried this, but since CuPy supports both `__cuda_array_interface__` and DLPack, you can most likely "launder" an array via CuPy into JAX:
* export the array via `__cuda_array_interface__` to CuPy.
* export the array via DLPack from CuPy.
* import the DLPack into JAX.

(Obviously this isn't ideal, but it might unblock you.)
Hi @hawkinsp I recently pm'd @apaszke in an occasion where this support was mentioned. It'd be nice if JAX can prioritize the bi-directional support for the CUDA Array Interface (and update to [the latest v3](https://numba.readthedocs.io/en/latest/cuda/cuda_array_interface.html) protocol, in which the synchronization semantics is specified). 

As you pointed out in a DLPack issue (https://github.com/dmlc/dlpack/issues/50), DLPack lacks the support for complex numbers and it's unlikely to be resolved in the foreseeable future. For array libraries this is simply not an acceptable workaround and is actually a blocker for several applications that I am aware. 

Thanks, and happy new year!
> Could this be reopened until import support is added as well?

Hi!

I was wondering if there is any update on this. Thanks!

Miguel
I'm very interested in this too

@hawkinsp this is seeing internal traction, given how some of the JAX internals have evolved (i.e. arrays, shmaps etc., in the context of MGMN) is there work being done here?
i would be interested in this
Just FYI, a new Python package [pydlpack](https://github.com/pearu/pydlpack) has been released that supports bidirectional data exchange of many array providers and consumers.

Here follows an example of zero-copy data exchange in between jax and torch:
```python
>>> from dlpack import asdlpack
>>> a1 = jax.numpy.array([[1, 2], [3, 4]])
>>> t1 = torch.from_dlpack(asdlpack(a1))
>>> t1
tensor([[1, 2],
        [3, 4]], device='cuda:0', dtype=torch.int32)
>>> t2 = torch.tensor([[5, 6], [7, 8]]).cuda()
>>> a2 = jax.numpy.from_dlpack(asdlpack(t2))
>>> a2
Array([[5, 6],
       [7, 8]], dtype=int32)
```

An other example is exchanging cuda buffers between jax and numba:
```python
>>> from dlpack import asdlpack
>>> import numba.cuda, numpy, jax
>>> a = numba.cuda.to_device(numpy.array([[1, 2], [3, 4]]))
>>> arr = jax.numpy.from_dlpack(asdlpack(a))
>>> arr
Array([[1, 2],
       [3, 4]], dtype=int32)
```

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -53,13 +53,14 @@
 from jax._src.custom_derivatives import custom_jvp
 from jax._src import dispatch
 from jax._src import dtypes
+from jax._src import xla_bridge
 from jax._src.api_util import _ensure_index_tuple
 from jax._src.array import ArrayImpl
 from jax._src.core import ShapedArray, ConcreteArray
 from jax._src.lax.lax import (_array_copy, _sort_lt_comparator,
                               _sort_le_comparator, PrecisionLike)
 from jax._src.lax import lax as lax_internal
-from jax._src.lib import xla_client as xc
+from jax._src.lib import xla_client as xc, xla_extension_version
 from jax._src.numpy import reductions
 from jax._src.numpy import ufuncs
 from jax._src.numpy import util
@@ -2121,9 +2122,25 @@ def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
   # to be used for type inference below.
   if isinstance(object, (bool, int, float, complex)):
     _ = dtypes.coerce_to_array(object, dtype)
+  elif not isinstance(object, Array):
+    # Check if object supports any of the data exchange protocols
+    # (except dlpack, see data-apis/array-api#301). If it does,
+    # consume the object as jax array and continue (but not return) so
+    # that other array() arguments get processed against the input
+    # object.
+    #
+    # Notice that data exchange protocols define dtype in the
+    # corresponding data structures and it may not be available as
+    # object.dtype. So, we'll resolve the protocols here before
+    # evaluating object.dtype.
+    if hasattr(object, '__jax_array__'):
+      object = object.__jax_array__()
+    elif hasattr(object, '__cuda_array_interface__'):
+      if xla_extension_version >= 237:
+        cai = object.__cuda_array_interface__
+        backend = xla_bridge.get_backend("cuda")
+        object = xc._xla.cuda_array_interface_to_buffer(cai, backend)
 
-  if hasattr(object, '__jax_array__'):
-    object = object.__jax_array__()
   object = tree_map(lambda leaf: leaf.__jax_array__()
                     if hasattr(leaf, "__jax_array__") else leaf, object)
   leaves = tree_leaves(object, is_leaf=lambda x: x is None)


Test patch:
diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py
--- a/tests/array_interoperability_test.py
+++ b/tests/array_interoperability_test.py
@@ -274,6 +274,87 @@ def testJaxToCuPy(self, shape, dtype):
                      z.__cuda_array_interface__["data"][0])
     self.assertAllClose(x, cupy.asnumpy(z))
 
+  @unittest.skipIf(xla_extension_version < 237, "Requires newer jaxlib")
+  @jtu.sample_product(
+    shape=all_shapes,
+    dtype=jtu.dtypes.supported(cuda_array_interface_dtypes),
+  )
+  @unittest.skipIf(not cupy, "Test requires CuPy")
+  @jtu.run_on_devices("cuda")
+  def testCuPyToJax(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    x = rng(shape, dtype)
+    y = cupy.asarray(x)
+    z = jnp.array(y, copy=False)  # this conversion uses dlpack protocol
+    self.assertEqual(z.dtype, dtype)
+    self.assertEqual(y.__cuda_array_interface__["data"][0],
+                     z.__cuda_array_interface__["data"][0])
+    self.assertAllClose(np.asarray(z), cupy.asnumpy(y))
+
+  @unittest.skipIf(xla_extension_version < 237, "Requires newer jaxlib")
+  @jtu.sample_product(
+    shape=all_shapes,
+    dtype=jtu.dtypes.supported(cuda_array_interface_dtypes),
+  )
+  @jtu.run_on_devices("cuda")
+  def testCaiToJax(self, shape, dtype):
+    if xb.using_pjrt_c_api():
+      self.skipTest("CUDA Array Interface support is incomplete in the PJRT C API")  # TODO(jakevdp)
+    rng = jtu.rand_default(self.rng())
+    x = rng(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]
+    with jax.default_device(device):
+      y = jnp.array(x, dtype=dtype)
+    self.assertEqual(y.dtype, dtype)
+
+    # Using a jax array CAI provider support to construct an object
+    # that implements the CUDA Array Interface, versions 2 and 3.
+    cai = y.__cuda_array_interface__
+    stream = tuple(y.devices())[0].get_stream_for_external_ready_events()
+
+    class CAIWithoutStridesV2:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 2
+      # CAI version 2 may not define strides and does not define stream
+      __cuda_array_interface__.pop("strides", None)
+      __cuda_array_interface__.pop("stream", None)
+
+    class CAIWithoutStrides:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 3
+      __cuda_array_interface__["strides"] = None
+      __cuda_array_interface__["stream"] = None  # default stream
+
+    class CAIWithStrides:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 3
+      strides = (dtype.dtype.itemsize,) if shape else ()
+      for s in reversed(shape[1:]):
+        strides = (strides[0] * s, *strides)
+      __cuda_array_interface__['strides'] = strides
+      __cuda_array_interface__["stream"] = stream
+
+    for CAIObject in [CAIWithoutStridesV2, CAIWithoutStrides,
+                      CAIWithStrides]:
+      z = jnp.array(CAIObject(), copy=False)
+      self.assertEqual(y.__cuda_array_interface__["data"][0],
+                       z.__cuda_array_interface__["data"][0])
+      self.assertAllClose(x, z)
+      if 0 in shape:
+        # the device id detection from a zero pointer value is not
+        # possible
+        pass
+      else:
+        self.assertEqual(y.devices(), z.devices())
+
+      z = jnp.array(CAIObject(), copy=True)
+      if 0 not in shape:
+        self.assertNotEqual(y.__cuda_array_interface__["data"][0],
+                            z.__cuda_array_interface__["data"][0])
+      self.assertAllClose(x, z)
 
 class Bfloat16Test(jtu.JaxTestCase):