def kmeans(
    observations: Array,
    centroids: Array,
    n_splits: int,
    threshold: float = 1e-5,
    max_iterations: Optional[int] = None,
    parallel_computation: bool = False) -> Tuple[Array, Array, float, int]:
  """Perform K-means clustering.

  Args:
    observations: [n_observations, dim] data points.
    centroids: [n_clusters, dim] initial cluster centroids.
    n_splits: split observations into this many chunks during scoring.
    threshold: stop computation if updates lead to improvement less than this.
    max_iterations: stop computation after this many iterations.
    parallel_computation: if true, assumes is run inside pmap.

  Returns:
    centroids: [n_clusters, dim] new cluster centroids.
    assignments: [n_observations] final assignment of observations to clusters.
    dist_diff: improvement in last k-means step.
    step: idx of next step.
  """
  # Run one step of K-means to get the initial value for the while loop.
  partial_step = functools.partial(
      kmeans_step, n_splits=n_splits, parallel_computation=parallel_computation)
  initial_val = partial_step((centroids, observations, jnp.inf, None, 0))

  def condition_fun(val) -> bool:
    below_max_iters = (max_iterations is None or val[4] < max_iterations)
    above_threshold = val[3] - val[2] > threshold
    should_continue = below_max_iters & above_threshold
    return should_continue

  centroids, _, mean_dist, prev_dist, step = jax.lax.while_loop(
      condition_fun, partial_step, initial_val)

  assignments, _ = compute_assignments(centroids, observations, n_splits)

  return centroids, assignments, prev_dist - mean_dist, step
Perform K-means clustering.

  Args:
    observations: [n_observations, dim] data points.
    centroids: [n_clusters, dim] initial cluster centroids.
    n_splits: split observations into this many chunks during scoring.
    threshold: stop computation if updates lead to improvement less than this.
    max_iterations: stop computation after this many iterations.
    parallel_computation: if true, assumes is run inside pmap.

  Returns:
    centroids: [n_clusters, dim] new cluster centroids.
    assignments: [n_observations] final assignment of observations to clusters.
    dist_diff: improvement in last k-means step.
    step: idx of next step.
  
language/mentionmemory/modules/kmeans.py