def _get_kmeans_lookup_table_and_weight(
    nbits, weight, force_kmeans1d=False, cluster_dim: int = 1, vector_axis: _Optional[int] = None
):
    """
    Generate K-Means lookup table given weights

    nbits:
        Number of bits for quantization

    weight:
        Weights as numpy array

    force_kmeans1d:
        Use kmeans1d regardless of number of weights

    Returns
    -------
    lut: numpy.array
        Lookup table, numpy array of shape (1 << nbits, )
    wq: numpy.array
        Quantized weight
    """
    if force_kmeans1d and cluster_dim > 1:
        raise ValueError("Cannot force kmeans1d for vector palettization (cluster_dim > 1).")

    num_weights = _np.prod(weight.shape)
    lut_len = 1 << nbits

    if cluster_dim > 1:
        weight = _optimize_utils.reshape_weight_for_vector_lut(weight, cluster_dim, vector_axis)

    weight = weight.reshape(-1, cluster_dim)
    lut = _np.zeros((lut_len, cluster_dim))

    is_better_to_use_kmeans1d = (
        weight.shape[1] == 1 and num_weights >= 10_000 and weight.dtype == _np.float16
    )

    if (is_better_to_use_kmeans1d and _HAS_KMEANS1D) or force_kmeans1d:
        # Cluster with kmeans1d
        assert _HAS_KMEANS1D, "Unable to import kmeans1d, please make sure it's installed."
        values, indices, counts = _np.unique(weight, return_inverse=True, return_counts=True)
        indices = indices.flatten()
        n_clusters = min(len(values), lut_len)
        kmeans_results = _kmeans1d.cluster(values, n_clusters, weights=counts)
        lut = lut.squeeze(-1)
        lut[:n_clusters] = kmeans_results.centroids
        wq = _np.array(kmeans_results.clusters)[indices]
    else:
        # Cluster with scikit-learn
        try:
            from sklearn.cluster import KMeans
        except:
            raise ModuleNotFoundError(
                "scikit-learn is required for k-means quantization."
                " To install, run: \"pip install scikit-learn\"."
            )

        if is_better_to_use_kmeans1d:
            _logger.warning("It would be better to use kmeans1d but that is not available."
                         " Using scikit-learn for K-means.")

        n_clusters = min(num_weights, lut_len)
        kmeans = KMeans(n_clusters, init="k-means++", tol=1e-2, n_init=1, random_state=0).fit(
            weight
        )
        wq = kmeans.labels_[:num_weights]
        lut[:n_clusters] = kmeans.cluster_centers_

    return lut, wq

    Generate K-Means lookup table given weights

    nbits:
        Number of bits for quantization

    weight:
        Weights as numpy array

    force_kmeans1d:
        Use kmeans1d regardless of number of weights

    Returns
    -------
    lut: numpy.array
        Lookup table, numpy array of shape (1 << nbits, )
    wq: numpy.array
        Quantized weight
    
coremltools/models/neural_network/quantization_utils.py