From ac0dea2e076ddff61368542806fe8695c94c3102 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 15 May 2023 15:00:41 +0100 Subject: [PATCH] Avoid DeprecationWarning in cupy dispatch registration In cupy 12 and above, cusparse should be imported from cupyx, not cupy itself. --- distributed/protocol/cupy.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/distributed/protocol/cupy.py b/distributed/protocol/cupy.py index dfc19a34128..fc6e0d08712 100644 --- 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