function updateCentroids(samples, nearestIndices, nClusters) {
  return tf.tidy(() => {
    const range = tf.range(0, nClusters).toInt();
    // mask.shape = [nSamples, nClusters]
    const mask = tf.equal(
      tf.expandDims(range, 0),
      tf.expandDims(nearestIndices, 1)
    );

    const expandedSamples = tf.expandDims(samples, 2);
    const expandedMask = tf.expandDims(mask, 1);
    const centroids = tf.div(
      // set all masked instances to 0 by multiplying the mask tensor, then sum across all instances
      tf.sum(tf.mul(expandedMask, expandedSamples), 0),
      // divided by number of instances
      tf.sum(mask, 0).toFloat()
    );
    // return.shape = [nClusters, nDims]
    return centroids.transpose();
  });
}
Calculate the average value of all instances in each cluster.
@param {Tensor2D} samples - data input, shape = [nInstances, nDims]
@param {Tensor2D} nearestIndices - indices of centroid to each data instance, shape = [nInstances, nClusters]
@param {Number} nClusters
@returns {Tensor2D} centroids of all instances in each cluster, shape = [nClusters, nDims]