ak.flatten and ak.ravel should test for nplike.ndarray, not np.ndarra…

…y. (#1340)
scikit-hep · Mar 3, 2022 · 62e15fd · 62e15fd
1 parent d4b04ae
commit 62e15fd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/awkward/_v2/operations/structure/ak_flatten.py b/src/awkward/_v2/operations/structure/ak_flatten.py
@@ -98,8 +98,9 @@ def flatten(array, axis=1, highlevel=True, behavior=None):
 
     if axis is None:
         out = layout.completely_flatten(function_name="ak.flatten")
-        assert isinstance(out, tuple) and all(isinstance(x, np.ndarray) for x in out)
-
+        assert isinstance(out, tuple) and all(
+            isinstance(x, nplike.ndarray) for x in out
+        )
         out = ak._v2.contents.NumpyArray(nplike.concatenate(out))
 
     elif axis == 0 or layout.axis_wrap_if_negative(axis) == 0:

diff --git a/src/awkward/_v2/operations/structure/ak_ravel.py b/src/awkward/_v2/operations/structure/ak_ravel.py
@@ -48,7 +48,7 @@ def ravel(array, highlevel=True, behavior=None):
     nplike = ak.nplike.of(layout)
 
     out = layout.completely_flatten(function_name="ak.ravel")
-    assert isinstance(out, tuple) and all(isinstance(x, np.ndarray) for x in out)
+    assert isinstance(out, tuple) and all(isinstance(x, nplike.ndarray) for x in out)
 
     if any(isinstance(x, nplike.ma.MaskedArray) for x in out):
         out = ak._v2.contents.NumpyArray(nplike.ma.concatenate(out))