Fix cupy.cublas import

Reported in #920.
explosion · Feb 7, 2024 · b3f4f47 · b3f4f47
1 parent b183006
commit b3f4f47
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/thinc/backends/cupy_ops.py b/thinc/backends/cupy_ops.py
@@ -1,7 +1,7 @@
 import numpy
 
 from .. import registry
-from ..compat import cupy, cupyx
+from ..compat import cublas, cupy, cupyx
 from ..types import DeviceTypes
 from ..util import (
     is_cupy_array,
@@ -257,7 +257,7 @@ def clip_gradient(self, gradient, threshold):
         # implementation.
         def frobenius_norm(X):
             X_vec = X.reshape(-1)
-            return cupy.cublas.nrm2(X_vec)
+            return cublas.nrm2(X_vec)
 
         grad_norm = cupy.maximum(frobenius_norm(gradient), 1e-12)
         gradient *= cupy.minimum(threshold, grad_norm) / grad_norm

diff --git a/thinc/compat.py b/thinc/compat.py
@@ -4,9 +4,11 @@
 
 try:  # pragma: no cover
     import cupy
+    import cupy.cublas
     import cupyx
 
     has_cupy = True
+    cublas = cupy.cublas
     cupy_version = Version(cupy.__version__)
     try:
         cupy.cuda.runtime.getDeviceCount()
@@ -20,6 +22,7 @@
     else:
         cupy_from_dlpack = cupy.fromDlpack
 except (ImportError, AttributeError):
+    cublas = None
     cupy = None
     cupyx = None
     cupy_version = Version("0.0.0")