def compute_assignments(
    centroids: Array,
    observations: Array,
    n_splits: int,
) -> Tuple[Array, Array]:
  """Assigns observations to cluster centroids.

  Computes l2 distance between each pair of observation and centroids, and
  assigns observation to closest cluster. Because the array of pairwise
  distances is very large, the cluster assignment is performed chunk by chunk.

  Args:
    centroids: [n_clusters, dim] cluster centroids.
    observations: [n_observations, dim] data points.
    n_splits: split observations into this many chunks.

  Returns:
    assignments: [n_observations] closest cluster for each observation.
    min_dist: [n_observations] distance to closest cluster by observation.
  """
  reshaped_observations = observations.reshape(n_splits, -1,
                                               centroids.shape[-1])

  def compute_split_assignments(split_points):
    dist = l2_distance(split_points, centroids)
    split_assignments = jnp.argmin(dist, axis=-1)
    split_min_dist = jnp.min(dist, axis=-1)
    return split_assignments, split_min_dist

  assignments, min_dist = jax.lax.map(compute_split_assignments,
                                      reshaped_observations)
  assignments = assignments.reshape(-1)
  min_dist = min_dist.reshape(-1)

  return assignments, min_dist
Assigns observations to cluster centroids.

  Computes l2 distance between each pair of observation and centroids, and
  assigns observation to closest cluster. Because the array of pairwise
  distances is very large, the cluster assignment is performed chunk by chunk.

  Args:
    centroids: [n_clusters, dim] cluster centroids.
    observations: [n_observations, dim] data points.
    n_splits: split observations into this many chunks.

  Returns:
    assignments: [n_observations] closest cluster for each observation.
    min_dist: [n_observations] distance to closest cluster by observation.
  
language/mentionmemory/modules/kmeans.py