prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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`.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `examples/kmeans/kmeans.py`
  Content:
  ```
  1 import argparse
  2 import contextlib
  3 import time
  4 
  5 import matplotlib.pyplot as plt
  6 import numpy as np
  7 import six
  8 
  9 import cupy
  10 
  11 
  12 @contextlib.contextmanager
  13 def timer(message):
  14     cupy.cuda.Stream.null.synchronize()
  15     start = time.time()
  16     yield
  17     cupy.cuda.Stream.null.synchronize()
  18     end = time.time()
  19     print('%s:  %f sec' % (message, end - start))
  20 
  21 
  22 _fit_calc_distances = cupy.ElementwiseKernel(
  23     'S data, raw S centers, int32 n_clusters, int32 dim', 'raw S dist',
  24     '''
  25     for (int j = 0; j < n_clusters; j++){
  26         int cent_ind[] = {j, i % dim};
  27         int dist_ind[] = {i / dim, j};
  28         double diff = centers[cent_ind] - data;
  29         atomicAdd(&dist[dist_ind], diff * diff);
  30     }
  31     ''',
  32     'calc_distances'
  33 )
  34 _fit_calc_center = cupy.ElementwiseKernel(
  35     'S data, T label, int32 dim', 'raw S centers, raw S group',
  36     '''
  37     int cent_ind[] = {label, i % dim};
  38     atomicAdd(&centers[cent_ind], data);
  39     atomicAdd(&group[label], 1);
  40     ''',
  41     'calc_center'
  42 )
  43 
  44 
  45 def fit(X, n_clusters, max_iter, use_custom_kernel):
  46     assert X.ndim == 2
  47     xp = cupy.get_array_module(X)
  48     pred = xp.zeros(len(X), dtype=np.int32)
  49     initial_indexes = np.random.choice(len(X), n_clusters,
  50                                        replace=False).astype(np.int32)
  51     centers = X[initial_indexes]
  52     data_num = X.shape[0]
  53     data_dim = X.shape[1]
  54 
  55     for _ in six.moves.range(max_iter):
  56         # calculate distances and label
  57         if not use_custom_kernel or xp == np:
  58             distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :],
  59                                        axis=2)
  60         else:
  61             distances = xp.zeros((data_num, n_clusters), dtype=np.float32)
  62             _fit_calc_distances(X, centers, n_clusters, data_dim, distances)
  63 
  64         new_pred = xp.argmin(distances, axis=1).astype(np.int32)
  65         if xp.all(new_pred == pred):
  66             break
  67         pred = new_pred
  68 
  69         # calculate centers
  70         if not use_custom_kernel or xp == np:
  71             centers = xp.stack([X[pred == i].mean(axis=0)
  72                                 for i in six.moves.range(n_clusters)])
  73         else:
  74             centers = xp.zeros((n_clusters, data_dim),
  75                                dtype=np.float32)
  76             group = xp.zeros(n_clusters, dtype=np.float32)
  77             label = pred[:, None]
  78             _fit_calc_center(X, label, data_dim, centers, group)
  79             group /= data_dim
  80             centers /= group[:, None]
  81 
  82     return centers, pred
  83 
  84 
  85 def draw(X, n_clusters, centers, pred, output):
  86     xp = cupy.get_array_module(X)
  87     for i in six.moves.range(n_clusters):
  88         labels = X[pred == i]
  89         if xp == cupy:
  90             labels = labels.get()
  91         plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(3))
  92     if xp == cupy:
  93         centers = centers.get()
  94     plt.scatter(centers[:, 0], centers[:, 1], s=120, marker='s',
  95                 facecolors='y', edgecolors='k')
  96     plt.savefig(output)
  97 
  98 
  99 def run(gpuid, n_clusters, num, max_iter, use_custom_kernel, output):
  100     samples = np.random.randn(num, 2).astype(np.float32)
  101     X_train = np.r_[samples + 1, samples - 1]
  102     repeat = 1
  103 
  104     with timer(' CPU '):
  105         for i in range(repeat):
  106             centers, pred = fit(X_train, n_clusters, max_iter,
  107                                 use_custom_kernel)
  108 
  109     with cupy.cuda.Device(gpuid):
  110         X_train = cupy.asarray(X_train)
  111         with timer(' GPU '):
  112             for i in range(repeat):
  113                 centers, pred = fit(X_train, n_clusters, max_iter,
  114                                     use_custom_kernel)
  115         if output is not None:
  116             index = np.random.choice(10000000, 300, replace=False)
  117             draw(X_train[index], n_clusters, centers, pred[index], output)
  118 
  119 
  120 if __name__ == '__main__':
  121     parser = argparse.ArgumentParser()
  122     parser.add_argument('--gpu-id', '-g', default=0, type=int,
  123                         help='ID of GPU.')
  124     parser.add_argument('--n-clusters', '-n', default=2, type=int,
  125                         help='number of clusters')
  126     parser.add_argument('--num', default=5000000, type=int,
  127                         help='number of samples')
  128     parser.add_argument('--max-iter', '-m', default=10, type=int,
  129                         help='number of iterations')
  130     parser.add_argument('--use-custom-kernel', action='store_true',
  131                         default=False, help='use Elementwise kernel')
  132     parser.add_argument('--output-image', '-o', default=None, type=str,
  133                         help='output image file name')
  134     args = parser.parse_args()
  135     run(args.gpu_id, args.n_clusters, args.num, args.max_iter,
  136         args.use_custom_kernel, args.output_image)
  137
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
       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 @@
       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 @@
               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 @@
           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
  
verification_info: {"golden_diff": "diff --git a/examples/kmeans/kmeans.py b/examples/kmeans/kmeans.py\n--- a/examples/kmeans/kmeans.py\n+++ b/examples/kmeans/kmeans.py\n@@ -19,26 +19,22 @@\n     print('%s:  %f sec' % (message, end - start))\n \n \n-_fit_calc_distances = cupy.ElementwiseKernel(\n-    'S data, raw S centers, int32 n_clusters, int32 dim', 'raw S dist',\n-    '''\n-    for (int j = 0; j < n_clusters; j++){\n-        int cent_ind[] = {j, i % dim};\n-        int dist_ind[] = {i / dim, j};\n-        double diff = centers[cent_ind] - data;\n-        atomicAdd(&dist[dist_ind], diff * diff);\n-    }\n-    ''',\n-    'calc_distances'\n+var_kernel = cupy.ElementwiseKernel(\n+    'T x0, T x1, T c0, T c1', 'T out',\n+    'out = (x0 - c0) * (x0 - c0) + (x1 - c1) * (x1 - c1)',\n+    'var_kernel'\n )\n-_fit_calc_center = cupy.ElementwiseKernel(\n-    'S data, T label, int32 dim', 'raw S centers, raw S group',\n-    '''\n-    int cent_ind[] = {label, i % dim};\n-    atomicAdd(&centers[cent_ind], data);\n-    atomicAdd(&group[label], 1);\n-    ''',\n-    'calc_center'\n+sum_kernel = cupy.ReductionKernel(\n+    'T x, S mask', 'T out',\n+    'mask ? x : 0',\n+    'a + b', 'out = a', '0',\n+    'sum_kernel'\n+)\n+count_kernel = cupy.ReductionKernel(\n+    'T mask', 'float32 out',\n+    'mask ? 1.0 : 0.0',\n+    'a + b', 'out = a', '0.0',\n+    'count_kernel'\n )\n \n \n@@ -49,8 +45,6 @@\n     initial_indexes = np.random.choice(len(X), n_clusters,\n                                        replace=False).astype(np.int32)\n     centers = X[initial_indexes]\n-    data_num = X.shape[0]\n-    data_dim = X.shape[1]\n \n     for _ in six.moves.range(max_iter):\n         # calculate distances and label\n@@ -58,8 +52,8 @@\n             distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :],\n                                        axis=2)\n         else:\n-            distances = xp.zeros((data_num, n_clusters), dtype=np.float32)\n-            _fit_calc_distances(X, centers, n_clusters, data_dim, distances)\n+            distances = var_kernel(X[:, None, 0], X[:, None, 1],\n+                                   centers[None, :, 1], centers[None, :, 0])\n \n         new_pred = xp.argmin(distances, axis=1).astype(np.int32)\n         if xp.all(new_pred == pred):\n@@ -67,17 +61,16 @@\n         pred = new_pred\n \n         # calculate centers\n+        i = xp.arange(n_clusters)\n+        mask = pred == i[:, None]\n         if not use_custom_kernel or xp == np:\n-            centers = xp.stack([X[pred == i].mean(axis=0)\n-                                for i in six.moves.range(n_clusters)])\n+            sums = xp.where(mask[:, :, None], X, 0).sum(axis=1)\n+            counts = xp.count_nonzero(mask, axis=1)\n+            centers = sums / counts\n         else:\n-            centers = xp.zeros((n_clusters, data_dim),\n-                               dtype=np.float32)\n-            group = xp.zeros(n_clusters, dtype=np.float32)\n-            label = pred[:, None]\n-            _fit_calc_center(X, label, data_dim, centers, group)\n-            group /= data_dim\n-            centers /= group[:, None]\n+            sums = sum_kernel(X, mask[:, :, None], axis=1)\n+            counts = count_kernel(mask, axis=1)\n+            centers = sums / counts\n \n     return centers, pred\n", "issue": "Use `ReductionKernel` in k-means example\nThe 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`.\n", "before_files": [{"content": "import argparse\nimport contextlib\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport six\n\nimport cupy\n\n\n@contextlib.contextmanager\ndef timer(message):\n    cupy.cuda.Stream.null.synchronize()\n    start = time.time()\n    yield\n    cupy.cuda.Stream.null.synchronize()\n    end = time.time()\n    print('%s:  %f sec' % (message, end - start))\n\n\n_fit_calc_distances = cupy.ElementwiseKernel(\n    'S data, raw S centers, int32 n_clusters, int32 dim', 'raw S dist',\n    '''\n    for (int j = 0; j < n_clusters; j++){\n        int cent_ind[] = {j, i % dim};\n        int dist_ind[] = {i / dim, j};\n        double diff = centers[cent_ind] - data;\n        atomicAdd(&dist[dist_ind], diff * diff);\n    }\n    ''',\n    'calc_distances'\n)\n_fit_calc_center = cupy.ElementwiseKernel(\n    'S data, T label, int32 dim', 'raw S centers, raw S group',\n    '''\n    int cent_ind[] = {label, i % dim};\n    atomicAdd(&centers[cent_ind], data);\n    atomicAdd(&group[label], 1);\n    ''',\n    'calc_center'\n)\n\n\ndef fit(X, n_clusters, max_iter, use_custom_kernel):\n    assert X.ndim == 2\n    xp = cupy.get_array_module(X)\n    pred = xp.zeros(len(X), dtype=np.int32)\n    initial_indexes = np.random.choice(len(X), n_clusters,\n                                       replace=False).astype(np.int32)\n    centers = X[initial_indexes]\n    data_num = X.shape[0]\n    data_dim = X.shape[1]\n\n    for _ in six.moves.range(max_iter):\n        # calculate distances and label\n        if not use_custom_kernel or xp == np:\n            distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :],\n                                       axis=2)\n        else:\n            distances = xp.zeros((data_num, n_clusters), dtype=np.float32)\n            _fit_calc_distances(X, centers, n_clusters, data_dim, distances)\n\n        new_pred = xp.argmin(distances, axis=1).astype(np.int32)\n        if xp.all(new_pred == pred):\n            break\n        pred = new_pred\n\n        # calculate centers\n        if not use_custom_kernel or xp == np:\n            centers = xp.stack([X[pred == i].mean(axis=0)\n                                for i in six.moves.range(n_clusters)])\n        else:\n            centers = xp.zeros((n_clusters, data_dim),\n                               dtype=np.float32)\n            group = xp.zeros(n_clusters, dtype=np.float32)\n            label = pred[:, None]\n            _fit_calc_center(X, label, data_dim, centers, group)\n            group /= data_dim\n            centers /= group[:, None]\n\n    return centers, pred\n\n\ndef draw(X, n_clusters, centers, pred, output):\n    xp = cupy.get_array_module(X)\n    for i in six.moves.range(n_clusters):\n        labels = X[pred == i]\n        if xp == cupy:\n            labels = labels.get()\n        plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(3))\n    if xp == cupy:\n        centers = centers.get()\n    plt.scatter(centers[:, 0], centers[:, 1], s=120, marker='s',\n                facecolors='y', edgecolors='k')\n    plt.savefig(output)\n\n\ndef run(gpuid, n_clusters, num, max_iter, use_custom_kernel, output):\n    samples = np.random.randn(num, 2).astype(np.float32)\n    X_train = np.r_[samples + 1, samples - 1]\n    repeat = 1\n\n    with timer(' CPU '):\n        for i in range(repeat):\n            centers, pred = fit(X_train, n_clusters, max_iter,\n                                use_custom_kernel)\n\n    with cupy.cuda.Device(gpuid):\n        X_train = cupy.asarray(X_train)\n        with timer(' GPU '):\n            for i in range(repeat):\n                centers, pred = fit(X_train, n_clusters, max_iter,\n                                    use_custom_kernel)\n        if output is not None:\n            index = np.random.choice(10000000, 300, replace=False)\n            draw(X_train[index], n_clusters, centers, pred[index], output)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--gpu-id', '-g', default=0, type=int,\n                        help='ID of GPU.')\n    parser.add_argument('--n-clusters', '-n', default=2, type=int,\n                        help='number of clusters')\n    parser.add_argument('--num', default=5000000, type=int,\n                        help='number of samples')\n    parser.add_argument('--max-iter', '-m', default=10, type=int,\n                        help='number of iterations')\n    parser.add_argument('--use-custom-kernel', action='store_true',\n                        default=False, help='use Elementwise kernel')\n    parser.add_argument('--output-image', '-o', default=None, type=str,\n                        help='output image file name')\n    args = parser.parse_args()\n    run(args.gpu_id, args.n_clusters, args.num, args.max_iter,\n        args.use_custom_kernel, args.output_image)\n", "path": "examples/kmeans/kmeans.py"}], "after_files": [{"content": "import argparse\nimport contextlib\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport six\n\nimport cupy\n\n\n@contextlib.contextmanager\ndef timer(message):\n    cupy.cuda.Stream.null.synchronize()\n    start = time.time()\n    yield\n    cupy.cuda.Stream.null.synchronize()\n    end = time.time()\n    print('%s:  %f sec' % (message, end - start))\n\n\nvar_kernel = cupy.ElementwiseKernel(\n    'T x0, T x1, T c0, T c1', 'T out',\n    'out = (x0 - c0) * (x0 - c0) + (x1 - c1) * (x1 - c1)',\n    'var_kernel'\n)\nsum_kernel = cupy.ReductionKernel(\n    'T x, S mask', 'T out',\n    'mask ? x : 0',\n    'a + b', 'out = a', '0',\n    'sum_kernel'\n)\ncount_kernel = cupy.ReductionKernel(\n    'T mask', 'float32 out',\n    'mask ? 1.0 : 0.0',\n    'a + b', 'out = a', '0.0',\n    'count_kernel'\n)\n\n\ndef fit(X, n_clusters, max_iter, use_custom_kernel):\n    assert X.ndim == 2\n    xp = cupy.get_array_module(X)\n    pred = xp.zeros(len(X), dtype=np.int32)\n    initial_indexes = np.random.choice(len(X), n_clusters,\n                                       replace=False).astype(np.int32)\n    centers = X[initial_indexes]\n\n    for _ in six.moves.range(max_iter):\n        # calculate distances and label\n        if not use_custom_kernel or xp == np:\n            distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :],\n                                       axis=2)\n        else:\n            distances = var_kernel(X[:, None, 0], X[:, None, 1],\n                                   centers[None, :, 1], centers[None, :, 0])\n\n        new_pred = xp.argmin(distances, axis=1).astype(np.int32)\n        if xp.all(new_pred == pred):\n            break\n        pred = new_pred\n\n        # calculate centers\n        i = xp.arange(n_clusters)\n        mask = pred == i[:, None]\n        if not use_custom_kernel or xp == np:\n            sums = xp.where(mask[:, :, None], X, 0).sum(axis=1)\n            counts = xp.count_nonzero(mask, axis=1)\n            centers = sums / counts\n        else:\n            sums = sum_kernel(X, mask[:, :, None], axis=1)\n            counts = count_kernel(mask, axis=1)\n            centers = sums / counts\n\n    return centers, pred\n\n\ndef draw(X, n_clusters, centers, pred, output):\n    xp = cupy.get_array_module(X)\n    for i in six.moves.range(n_clusters):\n        labels = X[pred == i]\n        if xp == cupy:\n            labels = labels.get()\n        plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(3))\n    if xp == cupy:\n        centers = centers.get()\n    plt.scatter(centers[:, 0], centers[:, 1], s=120, marker='s',\n                facecolors='y', edgecolors='k')\n    plt.savefig(output)\n\n\ndef run(gpuid, n_clusters, num, max_iter, use_custom_kernel, output):\n    samples = np.random.randn(num, 2).astype(np.float32)\n    X_train = np.r_[samples + 1, samples - 1]\n    repeat = 1\n\n    with timer(' CPU '):\n        for i in range(repeat):\n            centers, pred = fit(X_train, n_clusters, max_iter,\n                                use_custom_kernel)\n\n    with cupy.cuda.Device(gpuid):\n        X_train = cupy.asarray(X_train)\n        with timer(' GPU '):\n            for i in range(repeat):\n                centers, pred = fit(X_train, n_clusters, max_iter,\n                                    use_custom_kernel)\n        if output is not None:\n            index = np.random.choice(10000000, 300, replace=False)\n            draw(X_train[index], n_clusters, centers, pred[index], output)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--gpu-id', '-g', default=0, type=int,\n                        help='ID of GPU.')\n    parser.add_argument('--n-clusters', '-n', default=2, type=int,\n                        help='number of clusters')\n    parser.add_argument('--num', default=5000000, type=int,\n                        help='number of samples')\n    parser.add_argument('--max-iter', '-m', default=10, type=int,\n                        help='number of iterations')\n    parser.add_argument('--use-custom-kernel', action='store_true',\n                        default=False, help='use Elementwise kernel')\n    parser.add_argument('--output-image', '-o', default=None, type=str,\n                        help='output image file name')\n    args = parser.parse_args()\n    run(args.gpu_id, args.n_clusters, args.num, args.max_iter,\n        args.use_custom_kernel, args.output_image)\n", "path": "examples/kmeans/kmeans.py"}]}