src/transformers/models/imagegpt/image_processing_imagegpt.py/color_quantize
def color_quantize(x, clusters):
    x = x.reshape(-1, 3)
    d = squared_euclidean_distance(x, clusters)
    return np.argmin(d, axis=1)