code: def fit_custom(X, n_clusters, max_iter):
      assert X.ndim == 2
  
      n_samples = len(X)
  
      pred = cupy.zeros(n_samples)
  
      initial_indexes = cupy.random.choice(
          n_samples, n_clusters, replace=False
      )
      centers = X[initial_indexes]
  
      for _ in range(max_iter):
          distances = var_kernel(
              X[:, None, 0],
              X[:, None, 1],
              centers[None, :, 1],
              centers[None, :, 0],
          )
          new_pred = cupy.argmin(distances, axis=1)
          if cupy.all(new_pred == pred):
              break
          pred = new_pred
  
          i = cupy.arange(n_clusters)
          mask = pred == i[:, None]
          sums = sum_kernel(X, mask[:, :, None], axis=1)
          counts = count_kernel(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 3 for calling a global function with example usage: fit_custom(X, n_clusters, max_iter) and returns: centers,  pred