[FIX] intp -> uintp for cupy

This will need to handle -ve fill value for count

Author: dcherian

flox/aggregations.py

@@ -252,8 +252,8 @@ def __repr__(self) -> str:
     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])
@@ -281,7 +281,7 @@ def _mean_finalize(sum_, count):
     combine=("sum", "sum"),
     finalize=_mean_finalize,
     fill_value=(0, 0),
-    dtypes=(None, np.intp),
+    dtypes=(None, np.uintp),
     final_dtype=np.floating,
 )
 nanmean = Aggregation(
@@ -290,7 +290,7 @@ def _mean_finalize(sum_, count):
     combine=("sum", "sum"),
     finalize=_mean_finalize,
     fill_value=(0, 0),
-    dtypes=(None, np.intp),
+    dtypes=(None, np.uintp),
     final_dtype=np.floating,
 )
 
@@ -315,7 +315,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(
@@ -325,7 +325,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(
@@ -335,7 +335,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(
@@ -345,7 +345,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,
 )
 
@@ -368,7 +368,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))]
 
@@ -398,8 +398,8 @@ def _pick_second(*x):
     fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 argmin = Aggregation(
@@ -411,8 +411,8 @@ def _pick_second(*x):
     fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 nanargmax = Aggregation(
@@ -424,8 +424,8 @@ def _pick_second(*x):
     fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
-    dtypes=(None, np.intp),
-    final_dtype=np.intp,
+    dtypes=(None, np.uintp),
+    final_dtype=np.uintp,
 )
 
 nanargmin = Aggregation(
@@ -437,8 +437,8 @@ def _pick_second(*x):
     fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
-    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)
@@ -574,8 +574,10 @@ 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,)
     else:
         agg.min_count = 0