Better checks everywhere

scikit-hep · Feb 15, 2024 · 76acca1 · 76acca1
1 parent 30ab88c
commit 76acca1
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/awkward/_connect/jax/reducers.py b/src/awkward/_connect/jax/reducers.py
@@ -178,7 +178,7 @@ def apply(
             return ak.contents.NumpyArray(
                 array.backend.nplike.asarray(result, dtype=array.dtype)
             )
-        elif array.dtype.type in (np.complex128, np.complex64):
+        elif np.issubdtype(array.dtype, np.complexfloating):
             return ak.contents.NumpyArray(result.view(array.dtype))
         else:
             return ak.contents.NumpyArray(result, backend=array.backend)
@@ -209,7 +209,7 @@ def apply(
             jax.ops.segment_sum(jax.numpy.log(array.data), parents.data)
         )
 
-        if array.dtype.type in (np.complex128, np.complex64):
+        if np.issubdtype(array.dtype, np.complexfloating):
             return ak.contents.NumpyArray(
                 result.view(array.dtype), backend=array.backend
             )
@@ -327,7 +327,7 @@ def apply(
         result = jax.ops.segment_min(array.data, parents.data)
         result = jax.numpy.minimum(result, self._min_initial(self.initial, array.dtype))
 
-        if array.dtype.type in (np.complex128, np.complex64):
+        if np.issubdtype(array.dtype, np.complexfloating):
             return ak.contents.NumpyArray(
                 array.backend.nplike.asarray(
                     result.view(array.dtype), dtype=array.dtype
@@ -388,7 +388,7 @@ def apply(
         result = jax.ops.segment_max(array.data, parents.data)
 
         result = jax.numpy.maximum(result, self._max_initial(self.initial, array.dtype))
-        if array.dtype.type in (np.complex128, np.complex64):
+        if np.issubdtype(array.dtype, np.complexfloating):
             return ak.contents.NumpyArray(
                 array.backend.nplike.asarray(
                     result.view(array.dtype), dtype=array.dtype