def _cluster_weights_2d(
        cls,
        block_weight: _torch.Tensor,
        block_importance: _torch.Tensor,
        config: KMeansConfig,
        rank: int,
    ) -> _Tuple[_torch.Tensor, _torch.Tensor]:
        """
        Cluster weights such that each centroid is a 2d vector, i.e., cluster_dim > 1.
        If axis = 0: vectors are chosen with elements along the output channel dimension.
        Example:
            weight = [
                       [1, 2, 3, 4],
                       [5, 6, 7, 8],
                    ]
            axis = 0
            ========
            clustering is done for the 4 points below:
                [
                  [1, 5], ---> point 1
                  [2, 6], ---> point 2
                  [3, 7], ---> point 3
                  [4, 8], ---> point 4
                ]
            axis = 1
            ========
            clustering is done for the 4 points below:
                [
                  [1, 2], ---> point 1
                  [3, 4], ---> point 2
                  [5, 6], ---> point 3
                  [7, 8], ---> point 4
                ]
        """
        num_clusters = 2**config.n_bits

        # Convert weight from N-D to 2-D.
        # Apply 2-D k-means clustering on 2-D weights.
        if config.axis == 0:
            # (C_out, C_in, H, W) -> (C_in, C_out, H, W)
            # (C_in, C_out, H, W) -> (C_in * H * W * C_out // cluster_dim, cluster_dim)
            weight_2d = block_weight.transpose(0, 1).reshape(-1, config.cluster_dim)
            importance_2d = (
                block_importance.transpose(0, 1)
                .reshape(-1, config.cluster_dim)
                .sum(dim=1, keepdim=True)
                if block_importance is not None
                else None
            )
        else:
            # (C_out, C_in, H, W) -> (C_in, C_out * H * W) -> (C_out * H * W, C_in)
            # (C_out * H * W, C_in) -> (C_out * H * W * C_in // cluster_dim, cluster_dim)
            weight_2d = block_weight.reshape(-1, config.cluster_dim)
            importance_2d = (
                block_importance.reshape(-1, config.cluster_dim).sum(dim=1, keepdim=True)
                if block_importance is not None
                else None
            )

        # Optionally move tensors to GPU
        if _torch.cuda.is_available():
            device_id = rank % _torch.cuda.device_count()
            weight_2d = weight_2d.to(f"cuda:{device_id}")
            importance_2d = (
                importance_2d.to(f"cuda:{device_id}") if importance_2d is not None else None
            )

        kmeans_results = _EfficientKMeans(
            n_clusters=num_clusters,
            init="kmeans++",
            n_init=5,
            max_iter=300,
        ).fit(weight_2d, sample_weight=importance_2d)

        weight_2d.cpu()
        if importance_2d is not None:
            importance_2d.cpu()

        return kmeans_results.cluster_centers_.cpu(), kmeans_results.labels_.cpu()

        Cluster weights such that each centroid is a 2d vector, i.e., cluster_dim > 1.
        If axis = 0: vectors are chosen with elements along the output channel dimension.
        Example:
            weight = [
                       [1, 2, 3, 4],
                       [5, 6, 7, 8],
                    ]
            axis = 0
            ========
            clustering is done for the 4 points below:
                [
                  [1, 5], ---> point 1
                  [2, 6], ---> point 2
                  [3, 7], ---> point 3
                  [4, 8], ---> point 4
                ]
            axis = 1
            ========
            clustering is done for the 4 points below:
                [
                  [1, 2], ---> point 1
                  [3, 4], ---> point 2
                  [5, 6], ---> point 3
                  [7, 8], ---> point 4
                ]
        
coremltools/optimize/torch/_utils/k_means.py