import torch


def initialize_cluster_labels(num_clusters,
                              img_dimensions,
                              device):
  """Initializes uniform cluster labels for an image.

  This function is used to initialize cluster labels that uniformly
  partition a 2-D image.

  Args:
    num_clusters: A list of 2 integers for number of clusters in y
      and x axes.
    img_dimensions: A list of 2 integers for image's y and x dimension.

  Returns:
    A 2-D long tensor with shape specified by img_dimension.
  """
  y_labels = torch.linspace(
      0, num_clusters[0] - 1, img_dimensions[0], device=device).round_().long()
  x_labels = torch.linspace(
      0, num_clusters[1] - 1, img_dimensions[1], device=device).round_().long()
  y_labels = y_labels.view(-1, 1)
  x_labels = x_labels.view(1, -1)
  labels = y_labels + (y_labels.max() + 1) * x_labels

  return labels