Repository: cupy/cupy

Issue:
Use `ReductionKernel` in k-means example
The k-mean example uses `ElementwiseKernel` as an example of a custom kernel, but its algorithm is essentially reduction rather than mapping. It would be better to use `ReductionKernel` in the k-means example and to provide another example that uses `ElementwiseKernel`.


Discussion:

Resolving patch:
diff --git a/examples/kmeans/kmeans.py b/examples/kmeans/kmeans.py
--- a/examples/kmeans/kmeans.py
+++ b/examples/kmeans/kmeans.py
@@ -19,26 +19,22 @@ def timer(message):
     print('%s:  %f sec' % (message, end - start))
 
 
-_fit_calc_distances = cupy.ElementwiseKernel(
-    'S data, raw S centers, int32 n_clusters, int32 dim', 'raw S dist',
-    '''
-    for (int j = 0; j < n_clusters; j++){
-        int cent_ind[] = {j, i % dim};
-        int dist_ind[] = {i / dim, j};
-        double diff = centers[cent_ind] - data;
-        atomicAdd(&dist[dist_ind], diff * diff);
-    }
-    ''',
-    'calc_distances'
+var_kernel = cupy.ElementwiseKernel(
+    'T x0, T x1, T c0, T c1', 'T out',
+    'out = (x0 - c0) * (x0 - c0) + (x1 - c1) * (x1 - c1)',
+    'var_kernel'
 )
-_fit_calc_center = cupy.ElementwiseKernel(
-    'S data, T label, int32 dim', 'raw S centers, raw S group',
-    '''
-    int cent_ind[] = {label, i % dim};
-    atomicAdd(&centers[cent_ind], data);
-    atomicAdd(&group[label], 1);
-    ''',
-    'calc_center'
+sum_kernel = cupy.ReductionKernel(
+    'T x, S mask', 'T out',
+    'mask ? x : 0',
+    'a + b', 'out = a', '0',
+    'sum_kernel'
+)
+count_kernel = cupy.ReductionKernel(
+    'T mask', 'float32 out',
+    'mask ? 1.0 : 0.0',
+    'a + b', 'out = a', '0.0',
+    'count_kernel'
 )
 
 
@@ -49,8 +45,6 @@ def fit(X, n_clusters, max_iter, use_custom_kernel):
     initial_indexes = np.random.choice(len(X), n_clusters,
                                        replace=False).astype(np.int32)
     centers = X[initial_indexes]
-    data_num = X.shape[0]
-    data_dim = X.shape[1]
 
     for _ in six.moves.range(max_iter):
         # calculate distances and label
@@ -58,8 +52,8 @@ def fit(X, n_clusters, max_iter, use_custom_kernel):
             distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :],
                                        axis=2)
         else:
-            distances = xp.zeros((data_num, n_clusters), dtype=np.float32)
-            _fit_calc_distances(X, centers, n_clusters, data_dim, distances)
+            distances = var_kernel(X[:, None, 0], X[:, None, 1],
+                                   centers[None, :, 1], centers[None, :, 0])
 
         new_pred = xp.argmin(distances, axis=1).astype(np.int32)
         if xp.all(new_pred == pred):
@@ -67,17 +61,16 @@ def fit(X, n_clusters, max_iter, use_custom_kernel):
         pred = new_pred
 
         # calculate centers
+        i = xp.arange(n_clusters)
+        mask = pred == i[:, None]
         if not use_custom_kernel or xp == np:
-            centers = xp.stack([X[pred == i].mean(axis=0)
-                                for i in six.moves.range(n_clusters)])
+            sums = xp.where(mask[:, :, None], X, 0).sum(axis=1)
+            counts = xp.count_nonzero(mask, axis=1)
+            centers = sums / counts
         else:
-            centers = xp.zeros((n_clusters, data_dim),
-                               dtype=np.float32)
-            group = xp.zeros(n_clusters, dtype=np.float32)
-            label = pred[:, None]
-            _fit_calc_center(X, label, data_dim, centers, group)
-            group /= data_dim
-            centers /= group[:, None]
+            sums = sum_kernel(X, mask[:, :, None], axis=1)
+            counts = count_kernel(mask, axis=1)
+            centers = sums / counts
 
     return centers, pred
 


Test patch: