Avoid DeprecationWarning in cupy dispatch registration

In cupy 12 and above, cusparse should be imported from cupyx, not cupy itself.
dask · May 15, 2023 · ac0dea2 · ac0dea2
1 parent 693ad9d
commit ac0dea2
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/distributed/protocol/cupy.py b/distributed/protocol/cupy.py
@@ -68,7 +68,12 @@ def dask_deserialize_cupy_ndarray(header, frames):
 
 
 try:
-    from cupy.cusparse import MatDescriptor
+    from packaging.version import Version
+
+    if Version(cupy.__version__) >= Version("12"):
+        from cupyx.cusparse import MatDescriptor
+    else:
+        from cupy.cusparse import MatDescriptor
     from cupyx.scipy.sparse import spmatrix
 except ImportError:
     MatDescriptor = None