[FIX] intp -> unitp for cupy

This will need to handle -ve fill value for count

Author: dcherian

flox/aggregations.py

@@ -228,8 +228,8 @@ def __repr__(self):
     combine="sum",
     fill_value=0,
     final_fill_value=0,
-    dtypes=np.intp,
-    final_dtype=np.intp,
+    dtypes=np.uintp,
+    final_dtype=np.uintp,
 )
 
 # note that the fill values are the result of np.func([np.nan, np.nan])
@@ -250,7 +250,7 @@ def __repr__(self):
     combine=("sum", "sum"),
     finalize=lambda sum_, count: sum_ / count,
     fill_value=(0, 0),
-    dtypes=(None, np.intp),
+    dtypes=(None, np.uintp),
     final_dtype=np.floating,
 )
 nanmean = Aggregation(
@@ -259,7 +259,7 @@ def __repr__(self):
     combine=("sum", "sum"),
     finalize=lambda sum_, count: sum_ / count,
     fill_value=(0, 0),
-    dtypes=(None, np.intp),
+    dtypes=(None, np.uintp),
     final_dtype=np.floating,
 )
 
@@ -283,7 +283,7 @@ def _std_finalize(sumsq, sum_, count, ddof=0):
     finalize=_var_finalize,
     fill_value=0,
     final_fill_value=np.nan,
-    dtypes=(None, None, np.intp),
+    dtypes=(None, None, np.uintp),
     final_dtype=np.floating,
 )
 nanvar = Aggregation(
@@ -293,7 +293,7 @@ def _std_finalize(sumsq, sum_, count, ddof=0):
     finalize=_var_finalize,
     fill_value=0,
     final_fill_value=np.nan,
-    dtypes=(None, None, np.intp),
+    dtypes=(None, None, np.uintp),
     final_dtype=np.floating,
 )
 std = Aggregation(
@@ -303,7 +303,7 @@ def _std_finalize(sumsq, sum_, count, ddof=0):
     finalize=_std_finalize,
     fill_value=0,
     final_fill_value=np.nan,
-    dtypes=(None, None, np.intp),
+    dtypes=(None, None, np.uintp),
     final_dtype=np.floating,
 )
 nanstd = Aggregation(
@@ -313,7 +313,7 @@ def _std_finalize(sumsq, sum_, count, ddof=0):
     finalize=_std_finalize,
     fill_value=0,
     final_fill_value=np.nan,
-    dtypes=(None, None, np.intp),
+    dtypes=(None, None, np.uintp),
     final_dtype=np.floating,
 )
 
@@ -336,7 +336,7 @@ def argreduce_preprocess(array, axis):
     assert len(axis) == 1
     axis = axis[0]
 
-    idx = dask.array.arange(array.shape[axis], chunks=array.chunks[axis], dtype=np.intp)
+    idx = dask.array.arange(array.shape[axis], chunks=array.chunks[axis], dtype=np.uintp)
     # broadcast (TODO: is this needed?)
     idx = idx[tuple(slice(None) if i == axis else np.newaxis for i in range(array.ndim))]
 
@@ -362,8 +362,8 @@ def _zip_index(array_, idx_):
     fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 argmin = Aggregation(
@@ -375,8 +375,8 @@ def _zip_index(array_, idx_):
     fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 nanargmax = Aggregation(
@@ -388,8 +388,8 @@ def _zip_index(array_, idx_):
     fill_value=(dtypes.NINF, -1),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 nanargmin = Aggregation(
@@ -401,8 +401,8 @@ def _zip_index(array_, idx_):
     fill_value=(dtypes.INF, -1),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 first = Aggregation("first", chunk=None, combine=None, fill_value=0)
@@ -520,7 +520,9 @@ def _initialize_aggregation(
         agg.combine += ("sum",)
         agg.fill_value["intermediate"] += (0,)
         agg.fill_value["numpy"] += (0,)
-        agg.dtype["intermediate"] += (np.intp,)
-        agg.dtype["numpy"] += (np.intp,)
+        # uintp is supported by cupy, intp is not
+        # Also count is >=0, so uint should be fine.
+        agg.dtype["intermediate"] += (np.uintp,)
+        agg.dtype["numpy"] += (np.uintp,)
 
     return agg