Repository: cupy/cupy

Issue:
Windows: Fix csc/csr argmax/argmin tests
CuPy always returns in `int64`, but on Windows the default int type is `int32`.

```
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_2_{axis=-2, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_2_{axis=-2, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_3_{axis=-2, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_3_{axis=-2, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_4_{axis=-1, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_4_{axis=-1, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_5_{axis=-1, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_5_{axis=-1, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_6_{axis=0, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_6_{axis=0, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_7_{axis=0, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_7_{axis=0, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_8_{axis=1, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_8_{axis=1, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_9_{axis=1, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py::TestCscMatrixScipyCompressedMinMax_param_9_{axis=1, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_2_{axis=-2, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_2_{axis=-2, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_3_{axis=-2, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_3_{axis=-2, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_4_{axis=-1, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_4_{axis=-1, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_5_{axis=-1, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_5_{axis=-1, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_6_{axis=0, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_6_{axis=0, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_7_{axis=0, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_7_{axis=0, dense=True}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_8_{axis=1, dense=False}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_8_{axis=1, dense=False}::test_argmin
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_9_{axis=1, dense=True}::test_argmax
FAILED tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py::TestCsrMatrixScipyCompressedMinMax_param_9_{axis=1, dense=True}::test_argmin
```


Discussion:
Jenkins, test this please
Jenkins CI test (for commit 96be6e1a4a03183bfc4f68584f7c64fa6c996b1a, target branch master) succeeded!

/test
Windows failures are fixed!
On a second thought, perhaps instead of fixing the tests I should fix the return type of argmin/argmax to `'l'` so that it picks up the platform dependency? @kmaehashi WDYT?

Resolving patch:
diff --git a/cupyx/scipy/sparse/compressed.py b/cupyx/scipy/sparse/compressed.py
--- a/cupyx/scipy/sparse/compressed.py
+++ b/cupyx/scipy/sparse/compressed.py
@@ -1,17 +1,18 @@
-import numpy
+import string
+import warnings
 
+import numpy
 try:
     import scipy.sparse
     scipy_available = True
 except ImportError:
     scipy_available = False
 
-import warnings
-
 import cupy
 import cupyx
 
 from cupy import core
+from cupy.core import _scalar
 from cupy._creation import basic
 from cupy import cusparse
 from cupyx.scipy.sparse import base
@@ -175,66 +176,15 @@ class _compressed_sparse_matrix(sparse_data._data_matrix,
         }
         ''', 'min_nonzero_reduction')
 
-    _max_arg_reduction_kern = core.RawKernel(r'''
-        extern "C" __global__
-        void max_arg_reduction(double* data, int* indices, int* x, int* y,
-                               int length, long long* z) {
-            // Get the index of the block
-            int tid = blockIdx.x * blockDim.x + threadIdx.x;
-
-            // Calculate the block length
-            int block_length = y[tid] - x[tid];
-
-            // Select initial value based on the block density
-            int data_index = 0;
-            double data_value = 0;
-
-            if (block_length == length){
-                // Block is dense. Fill the first value
-                data_value = data[x[tid]];
-                data_index = indices[x[tid]];
-            } else if (block_length > 0)  {
-                // Block has at least one zero. Assign first occurrence as the
-                // starting reference
-                data_value = 0;
-                for (data_index = 0; data_index < length; data_index++){
-                    if (data_index != indices[x[tid] + data_index] ||
-                        x[tid] + data_index >= y[tid]){
-                        break;
-                    }
-                }
-            } else {
-                 // Zero valued array
-                data_value = 0;
-                data_index = 0;
-            }
-
-            // Iterate over the section of the sparse matrix
-            for (int entry = x[tid]; entry < y[tid]; entry++){
-                if (data[entry] != data[entry]){
-                    // Check for NaN
-                    data_value = nan("");
-                    data_index = 0;
-                    break;
-                } else {
-                    // Check for a value update
-                    if (data[entry] > data_value){
-                        data_index = indices[entry];
-                        data_value = data[entry];
-                    }
-                }
-            }
-
-            // Store in the return function
-            z[tid] = data_index;
-        }
-        ''', 'max_arg_reduction')
+    # For _max_arg_reduction_mod and _min_arg_reduction_mod below, we pick
+    # the right template specialization according to input dtypes at runtime.
+    # The distinction in int types (T2) is important for portability in OS.
 
-    _min_arg_reduction_kern = core.RawKernel(r'''
-        extern "C" __global__
-        void min_arg_reduction(double* data, int* indices, int* x, int* y,
-                               int length, long long* z) {
-            // Get the index of hte block
+    _argmax_argmin_code = r'''
+        template<typename T1, typename T2> __global__ void
+        ${func}_arg_reduction(T1* data, int* indices, int* x, int* y,
+                              int length, T2* z) {
+            // Get the index of the block
             int tid = blockIdx.x * blockDim.x + threadIdx.x;
 
             // Calculate the block length
@@ -273,7 +223,7 @@ class _compressed_sparse_matrix(sparse_data._data_matrix,
                     break;
                 } else {
                     // Check for a value update
-                    if (data[entry] < data_value){
+                    if (data[entry] ${op} data_value){
                         data_index = indices[entry];
                         data_value = data[entry];
                     }
@@ -282,9 +232,25 @@ class _compressed_sparse_matrix(sparse_data._data_matrix,
 
             // Store in the return function
             z[tid] = data_index;
-
-        }
-        ''', 'min_arg_reduction')
+        }'''
+
+    _max_arg_reduction_mod = core.RawModule(
+        code=string.Template(_argmax_argmin_code).substitute(
+            func='max', op='>'),
+        options=('-std=c++11',),
+        name_expressions=['max_arg_reduction<float, int>',
+                          'max_arg_reduction<float, long long>',
+                          'max_arg_reduction<double, int>',
+                          'max_arg_reduction<double, long long>'])
+
+    _min_arg_reduction_mod = core.RawModule(
+        code=string.Template(_argmax_argmin_code).substitute(
+            func='min', op='<'),
+        options=('-std=c++11',),
+        name_expressions=['min_arg_reduction<float, int>',
+                          'min_arg_reduction<float, long long>',
+                          'min_arg_reduction<double, int>',
+                          'min_arg_reduction<double, long long>'])
 
     # TODO(leofang): rewrite a more load-balanced approach than this naive one?
     _has_sorted_indices_kern = core.ElementwiseKernel(
@@ -1026,42 +992,46 @@ def _arg_minor_reduce(self, ufunc, axis):
         # Call to the appropriate kernel function
         if axis == 1:
             # Create the vector to hold output
-            value = cupy.zeros(self.shape[0]).astype(cupy.int64)
+            # Note: it's important to set "int" here, following what SciPy
+            # does, as the outcome dtype is platform dependent
+            value = cupy.zeros(self.shape[0], dtype=int)
 
             # Perform the calculation
+            ker_name = '_arg_reduction<{}, {}>'.format(
+                _scalar.get_typename(self.data.dtype),
+                _scalar.get_typename(value.dtype))
             if ufunc == cupy.argmax:
-                self._max_arg_reduction_kern(
-                    (self.shape[0],), (1,),
-                    (self.data.astype(cupy.float64), self.indices,
-                     self.indptr[:len(self.indptr) - 1],
-                     self.indptr[1:], cupy.int64(self.shape[1]),
-                     value))
-            if ufunc == cupy.argmin:
-                self._min_arg_reduction_kern(
-                    (self.shape[0],), (1,),
-                    (self.data.astype(cupy.float64), self.indices,
-                     self.indptr[:len(self.indptr) - 1],
-                     self.indptr[1:], cupy.int64(self.shape[1]),
-                     value))
-
-        if axis == 0:
+                ker_name = 'max' + ker_name
+                ker = self._max_arg_reduction_mod.get_function(ker_name)
+            elif ufunc == cupy.argmin:
+                ker_name = 'min' + ker_name
+                ker = self._min_arg_reduction_mod.get_function(ker_name)
+            ker((self.shape[0],), (1,),
+                (self.data, self.indices,
+                 self.indptr[:len(self.indptr) - 1],
+                 self.indptr[1:], cupy.int64(self.shape[1]),
+                 value))
+
+        elif axis == 0:
             # Create the vector to hold output
-            value = cupy.zeros(self.shape[1]).astype(cupy.int64)
+            # Note: it's important to set "int" here, following what SciPy
+            # does, as the outcome dtype is platform dependent
+            value = cupy.zeros(self.shape[1], dtype=int)
 
             # Perform the calculation
+            ker_name = '_arg_reduction<{}, {}>'.format(
+                _scalar.get_typename(self.data.dtype),
+                _scalar.get_typename(value.dtype))
             if ufunc == cupy.argmax:
-                self._max_arg_reduction_kern(
-                    (self.shape[1],), (1,),
-                    (self.data.astype(cupy.float64), self.indices,
-                     self.indptr[:len(self.indptr) - 1],
-                     self.indptr[1:], cupy.int64(self.shape[0]),
-                     value))
-            if ufunc == cupy.argmin:
-                self._min_arg_reduction_kern(
-                    (self.shape[1],), (1,),
-                    (self.data.astype(cupy.float64), self.indices,
-                     self.indptr[:len(self.indptr) - 1],
-                     self.indptr[1:],
-                     cupy.int64(self.shape[0]), value))
+                ker_name = 'max' + ker_name
+                ker = self._max_arg_reduction_mod.get_function(ker_name)
+            elif ufunc == cupy.argmin:
+                ker_name = 'min' + ker_name
+                ker = self._min_arg_reduction_mod.get_function(ker_name)
+            ker((self.shape[1],), (1,),
+                (self.data, self.indices,
+                 self.indptr[:len(self.indptr) - 1],
+                 self.indptr[1:], cupy.int64(self.shape[0]),
+                 value))
 
         return value


Test patch: