def kmeans_step(
    val: Tuple[Array, Array, float, Optional[float], int],
    n_splits: int,
    parallel_computation: bool = False,
) -> Tuple[Array, Array, float, float, int]:
  """Perform single K-means step.

  Standard K-means step. Assigns observations to nearest cluster, then updates
  cluster centroids as mean of assigned observations. Inputs are packed into
  'val' to facilitate while loop condition for K-means.

  Args:
    val: tuple of [n_clusters, dim] cluster centroids. [n_observations, dim]
      observations. prev_dist, mean distance between observations and closest
      cluster centroid in previous K-means step. prev_2_dist, distance from two
      steps prior. step, idx of current step.
    n_splits: number of splits for compute assignments
    parallel_computation: if true, assumes is run inside pmap with
      'observations' axis.

  Returns:
    new_centroids: [n_clusters, dim] new cluster centroids.
    observations: [n_observations, dim].
    mean_dist: mean distance between observations and closest cluster centroid.
    prev_dist: mean distance for previous K-means step.
    step: idx of next step.
  """
  centroids, observations, prev_dist, _, step = val
  assignments, min_dist = compute_assignments(centroids, observations, n_splits)

  mean_dist = jnp.mean(min_dist)

  # Compute new cluster centroids as average of observations in cluster
  cluster_sums = jnp.zeros(centroids.shape).at[assignments].add(observations)
  counts = jnp.bincount(assignments, length=centroids.shape[0])

  if parallel_computation:
    cluster_sums = jax.lax.psum(cluster_sums, axis_name='observations')
    counts = jax.lax.psum(counts, axis_name='observations')
    mean_dist = jax.lax.pmean(mean_dist, axis_name='observations')

  new_centroids = cluster_sums / counts[:, None].clip(min=1.0)

  jax.debug.print('step={}', step)
  jax.debug.print('dist_diff={}', mean_dist - prev_dist)
  step += 1

  return new_centroids, observations, mean_dist, prev_dist, step  # type: ignore
Perform single K-means step.

  Standard K-means step. Assigns observations to nearest cluster, then updates
  cluster centroids as mean of assigned observations. Inputs are packed into
  'val' to facilitate while loop condition for K-means.

  Args:
    val: tuple of [n_clusters, dim] cluster centroids. [n_observations, dim]
      observations. prev_dist, mean distance between observations and closest
      cluster centroid in previous K-means step. prev_2_dist, distance from two
      steps prior. step, idx of current step.
    n_splits: number of splits for compute assignments
    parallel_computation: if true, assumes is run inside pmap with
      'observations' axis.

  Returns:
    new_centroids: [n_clusters, dim] new cluster centroids.
    observations: [n_observations, dim].
    mean_dist: mean distance between observations and closest cluster centroid.
    prev_dist: mean distance for previous K-means step.
    step: idx of next step.
  
language/mentionmemory/modules/kmeans.py