def _cluster_weights_with_masking(
        cls,
        block_weight: _torch.Tensor,
        block_importance: _torch.Tensor,
        block_mask: _torch.Tensor,
        config: KMeansConfig,
    ) -> _Tuple[_Optional[_torch.Tensor], _Optional[_torch.Tensor]]:
        """
        Cluster block weight with clustering only applied to masked weight elements.
        """
        num_clusters = 2**config.n_bits

        block_weight_flatten = block_weight.flatten()
        block_weight_flatten_masked = block_weight_flatten[block_mask]

        if len(block_weight_flatten_masked) > 0:
            if block_importance is not None:
                block_importance_flatten = block_importance.flatten()
                kmeans_results = _kmeans1d.cluster(
                    block_weight_flatten_masked.numpy(),
                    num_clusters,
                    weights=block_importance_flatten[block_mask].numpy(),
                )
            else:
                kmeans_results = _kmeans1d.cluster(
                    block_weight_flatten_masked.numpy(), num_clusters
                )
            return _torch.tensor(kmeans_results.centroids), _torch.tensor(kmeans_results.clusters)

        return None, None

        Cluster block weight with clustering only applied to masked weight elements.
        
coremltools/optimize/torch/_utils/k_means.py