code: def fit_xp(X, n_clusters, max_iter):
      assert X.ndim == 2
  
      # Get NumPy or CuPy module from the supplied array.
      xp = cupy.get_array_module(X)
  
      n_samples = len(X)
  
      # Make an array to store the labels indicating which cluster each sample is
      # contained.
      pred = xp.zeros(n_samples)
  
      # Choose the initial centroid for each cluster.
      initial_indexes = xp.random.choice(
          n_samples, n_clusters, replace=False
      )
      centers = X[initial_indexes]
  
      for _ in range(max_iter):
          # Compute the new label for each sample.
          distances = xp.linalg.norm(
              X[:, None, :] - centers[None, :, :], axis=2
          )
          new_pred = xp.argmin(distances, axis=1)
  
          # If the label is not changed for each sample, we suppose the
          # algorithm has converged and exit from the loop.
          if xp.all(new_pred == pred):
              break
          pred = new_pred
  
          # Compute the new centroid for each cluster.
          i = xp.arange(n_clusters)
          mask = pred == i[:, None]
          sums = xp.where(mask[:, :, None], X, 0).sum(axis=1)
          counts = xp.count_nonzero(mask, axis=1).reshape(
              (n_clusters, 1)
          )
          centers = sums / counts
  
      return centers, pred
  
file_path: cupy/examples/kmeans/kmeans.py
imports: import argparse
  import contextlib
  import time
  import cupy
  import matplotlib.pyplot
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: fit_xp(X, n_clusters, max_iter) and returns: centers,  pred