code: def KMeans(x, K=10, Niter=10000, verbose=True):
      """
      Implements Lloyd's algorithm for the Euclidean metric.
      """
      use_cuda = torch.cuda.is_available()
      dtype = torch.float32 if use_cuda else torch.float64
  
      start = time.time()
      (
          N,
          D,
      ) = x.shape  # Number of samples, dimension of the ambient space
      c = x[
          :K, :
      ].clone()  # Simplistic initialization for the centroids
  
      x_i = LazyTensor(x.view(N, 1, D))  # (N, 1, D) samples
      c_j = LazyTensor(c.view(1, K, D))  # (1, K, D) centroids
  
      # K-means loop:
      # - x  is the (N, D) point cloud,
      # - cl is the (N,) vector of class labels
      # - c  is the (K, D) cloud of cluster centroids
      for i in range(Niter):
          # E step: assign points to the closest cluster -------------------------
          D_ij = ((x_i - c_j) ** 2).sum(
              -1
          )  # (N, K) symbolic squared distances
          # print(D_ij.argmin(dim=1))
          # import pdb;pdb.set_trace()
          cl = (
              D_ij.argmin(dim=1).long().view(-1)
          )  # Points -> Nearest cluster
  
          # M step: update the centroids to the normalized cluster average: ------
          # Compute the sum of points per cluster:
          c.zero_()
          c.scatter_add_(0, cl[:, None].repeat(1, D), x)
  
          # Divide by the number of points per cluster:
          Ncl = torch.bincount(cl, minlength=K).type_as(c).view(K, 1)
          c /= Ncl  # in-place division to compute the average
  
      if (
          verbose
      ):  # Fancy display -----------------------------------------------
          if use_cuda:
              torch.cuda.synchronize()
          end = time.time()
          print(
              f"K-means for the Euclidean metric with {N:,} points in dimension"
              f" {D:,}, K = {K:,}:"
          )
          print(
              "Timing for {} iterations: {:.5f}s = {} x {:.5f}s\n".format(
                  Niter, end - start, Niter, (end - start) / Niter
              )
          )
  
      return cl, c
  
file_path: gluonts/src/gluonts/nursery/SCott/dataset_tools/algo_clustering.py
imports: from pts.dataset.repository.datasets import get_dataset, dataset_recipes
  from pts.dataset.utils import to_pandas
  import numpy
  import pandas
  from pts.dataset import ListDataset
  import torch
  from pts.model.simple_feedforward import SimpleFeedForwardEstimator
  import pickle
  import json
  import random
  import time
  from matplotlib import pyplot
  from pykeops.torch import LazyTensor
label_desc: Use image node_id 1 for calling a global function with example usage: KMeans(x, K, Niter, verbose) and returns: cl,  c