Don't accidentally iterate cupy with CPU

scikit-hep · Aug 13, 2024 · 946259c · 946259c
1 parent 8516135
commit 946259c
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 14 deletions.
diff --git a/src/awkward/contents/bitmaskedarray.py b/src/awkward/contents/bitmaskedarray.py
@@ -689,20 +689,17 @@ def _to_arrow(
         )
 
     def _to_cudf(self, cudf: Any, mask: Content | None, length: int):
-        cupy = Cupy.instance()
-        np = (
-            Numpy.instance()._module
-        )  # flip and resize are not in the arraylike instance
+        cp = Cupy.instance()._module
 
         assert mask is None  # this class has its own mask
         if not self.lsb_order:
-            m = np.flip(np.packbits(np.flip(np.unpackbits(self._mask.data))))
+            m = cp.flip(cp.packbits(cp.flip(cp.unpackbits(cp.asarray(self._mask.data)))))
         else:
             m = self._mask.data
 
         if m.nbytes % 64:
-            m = np.resize(m, ((m.nbytes // 64) + 1) * 64)
-        m = cudf.core.buffer.as_buffer(cupy.asarray(m))
+            m = cp.resize(m, ((m.nbytes // 64) + 1) * 64)
+        m = cudf.core.buffer.as_buffer(m)
         inner = self._content._to_cudf(cudf, mask=None, length=length)
         inner.set_base_mask(m)
         return inner

diff --git a/src/awkward/contents/bytemaskedarray.py b/src/awkward/contents/bytemaskedarray.py
@@ -1053,14 +1053,13 @@ def _to_arrow(
         )
 
     def _to_cudf(self, cudf: Any, mask: Content | None, length: int):
-        cupy = Cupy.instance()
-        np = Numpy.instance()._module
+        cp = Cupy.instance()._module
 
         assert mask is None  # this class has its own mask
-        m = np.packbits(self._mask, bitorder="little")
+        m = cp.packbits(cp.asarray(self._mask), bitorder="little")
         if m.nbytes % 64:
-            m.resize(((m.nbytes // 64) + 1) * 64)
-        m = cudf.core.buffer.as_buffer(cupy.asarray(m))
+            m = cp.resize(m, ((m.nbytes // 64) + 1) * 64)
+        m = cudf.core.buffer.as_buffer(m)
         inner = self._content._to_cudf(cudf, mask=None, length=length)
         inner.set_base_mask(m)
         return inner

diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py
@@ -1223,14 +1223,13 @@ def _to_arrow(
 
     def _to_cudf(self, cudf: Any, mask: Content | None, length: int):
         cupy = Cupy.instance()
-        np = Numpy.instance()
 
         assert self._backend.nplike.known_data
         data = cupy.asarray(self._data)
         if mask is not None:
             m = np.packbits(mask, bitorder="little")
             if m.nbytes % 64:
-                m.resize(((m.nbytes // 64) + 1) * 64)
+                m = cupy.resize(m, ((m.nbytes // 64) + 1) * 64)
             m = cudf.core.buffer.as_buffer(cupy.asarray(m))
         else:
             m = None