input: Write a python function to perform KMeans clustering and predict clusters for a given dataset.
  assert kmeans_clustering([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]], 2) in [[0, 0, 0, 1, 1, 1], [1, 1, 1, 0, 0, 0]]
  assert kmeans_clustering([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10, 0]], 2) in [[0, 0, 0, 1, 1, 1], [1, 1, 1, 0, 0, 0]]
  assert kmeans_clustering([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10, 0]], 3) in [[0, 0, 0, 1, 1, 2], [1, 1, 1, 0, 0, 2], [1, 1, 1, 2, 2, 0]]
output: def kmeans_clustering(data, k):
      from sklearn.cluster import KMeans
      kmeans = KMeans(n_clusters=k)
      kmeans.fit(data)
      return kmeans.labels_.tolist()