Function compute_partitions

pub fn compute_partitions<T: Float + L2 + Dot + Sync>(
    centroids: &[T],
    vectors: &[T],
    dimension: impl AsPrimitive<usize>,
    distance_type: DistanceType,
) -> Vec<Option<u32>>

Compute partition ID of each vector in the KMeans.

If returns None, means the vector is not valid, i.e., all NaN.