computeClusters = (
  data,
  nClusters,
  transposeInput = false,
  maxIterations = 300,
  tolerance = 10e-4
) => {
  if (!data.length) {
    return [];
  }
  return tf.tidy(() => {
    let samples = tf.tensor2d(data).toFloat();
    if (transposeInput) {
      samples = tf.transpose(samples);
    }
    let centroids = initCentroids(samples, nClusters);
    let nearest;
    for (let i = 0; i < maxIterations; i++) {
      nearest = assignClusterId(samples, centroids);
      centroids = updateCentroids(samples, nearest, nClusters);
    }
    // return an array of indices, representing the nearest "centroid" to each data point
    return nearest.dataSync();
  });
}
kmeans
@param {Array<Array>} data 2D data to run kmeans on. shape = [nInstances, nFeatures]
@param {Number} nClusters
@param {Boolean} transposeInput if the input data array shape is [nFeatures, nInstances], use transposeInput = true
@param {Number} maxIterations number of iterations before force stop
@param {Number} tolerance