def test_compress_cluster_weights_call(
    mocker, num_kmeans_workers, model, sensitivity_dict, config, kmeans_keys, request
):
    """
    Test ParallelKMeans/SequentialKMeans are called with correct arguments
    """
    if isinstance(model, str):
        model = request.getfixturevalue(model)
    if isinstance(sensitivity_dict, str):
        sensitivity_dict = request.getfixturevalue(sensitivity_dict)

    mocker.patch.object(SKMPalettizer, "compute_sensitivity", Mock(return_value=sensitivity_dict))
    mock_parallel = mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._ParallelKMeans"
    )
    mock_sequential = mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._SequentialKMeans"
    )

    palettizer = SKMPalettizer(model, config)

    palettizer.compress(
        num_sensitivity_workers=1,
        fsdp_auto_wrap_policy=None,
        num_kmeans_workers=num_kmeans_workers,
    )

    k_means_config_dict = {}
    for key, val in kmeans_keys.items():
        sensitivity_key = f"{key}.{val}" if len(key) > 0 else val
        k_means_config_dict[key] = {
            val: KMeansConfig(
                n_bits=ModuleSKMPalettizerConfig().n_bits,
                axis=0,
                block_size=None,
                cluster_dim=1,
                importance=sensitivity_dict[sensitivity_key],
                enable_per_channel_scale=ModuleSKMPalettizerConfig().enable_per_channel_scale,
            )
        }

    if num_kmeans_workers > 1:
        mock_parallel.cluster_weights.assert_called_once_with(
            palettizer._model, k_means_config_dict, num_workers=num_kmeans_workers
        )
    else:
        mock_sequential.cluster_weights.assert_called_once_with(
            palettizer._model,
            k_means_config_dict,
        )

    Test ParallelKMeans/SequentialKMeans are called with correct arguments
    
coremltools/test/optimize/torch/palettization/test_sensitive_k_means.py