def test_fsdp_auto_wrap_policy_compress_call(
    mocker, num_kmeans_workers, num_sensitivity_workers, auto_wrap_policy
):
    """
    Test compress passes fsdp_auto_wrap_policy argument correctly to
    compute_sensitivity method.
    """
    mock_compute_sensitivity = Mock(return_value={"weight": None})

    mocker.patch.object(SKMPalettizer, "compute_sensitivity", mock_compute_sensitivity)
    mocker.patch("coremltools.optimize.torch.palettization.sensitive_k_means._ParallelKMeans")
    mocker.patch("coremltools.optimize.torch.palettization.sensitive_k_means._SequentialKMeans")

    model = torch.nn.Linear(5, 10)
    palettizer = SKMPalettizer(model)

    palettizer.compress(
        num_sensitivity_workers=num_sensitivity_workers,
        fsdp_auto_wrap_policy=auto_wrap_policy,
        num_kmeans_workers=num_kmeans_workers,
    )

    mock_compute_sensitivity.assert_called_once_with(
        None,
        None,
        None,
        num_sensitivity_workers,
        fsdp_auto_wrap_policy=auto_wrap_policy,
    )

    Test compress passes fsdp_auto_wrap_policy argument correctly to
    compute_sensitivity method.
    
coremltools/test/optimize/torch/palettization/test_sensitive_k_means.py