def test_fsdp_auto_wrap_policy_multi_worker_compute_sensitivity_call(mocker, auto_wrap_policy):
    """
    Test _compute_sensitivity_impl_multiple_workers passes correct value of fsdp auto wrap policy
    to FSDP call
    """
    model = torch.nn.Linear(5, 10)

    mocker.patch("coremltools.optimize.torch.palettization.sensitive_k_means._torch")
    mocker.patch("coremltools.optimize.torch.palettization.sensitive_k_means._ddp_setup")
    mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._is_leader",
        Mock(return_value=True),
    )
    mocker.patch.object(
        SKMPalettizer, "_register_grad_square_hooks", Mock(return_value=nullcontext())
    )

    if auto_wrap_policy is not None:
        expected_auto_wrap_policy = Mock()
        auto_wrap_policy.get_policy.return_value = expected_auto_wrap_policy
    else:
        expected_auto_wrap_policy = None

    with patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._FSDP", autospec=True
    ) as mock_fsdp:
        mock_fsdp.state_dict_type.return_value = nullcontext()

        palettizer = SKMPalettizer(model)

        palettizer._compute_sensitivity_impl_multiple_workers(
            rank=0,
            num_workers=1,
            dataset=[None],
            loss_fn=Mock(),
            sensitivity_path=None,
            fsdp_auto_wrap_policy=auto_wrap_policy,
        )

        # test FSDP either gets None or output of get_policy method on the
        # FSDPAutoWrapPolicy object
        mock_fsdp.assert_called_with(
            module=palettizer._model,
            auto_wrap_policy=expected_auto_wrap_policy,
            sharding_strategy=ANY,
            use_orig_params=False,
            device_id=ANY,
            sync_module_states=True,
        )

    Test _compute_sensitivity_impl_multiple_workers passes correct value of fsdp auto wrap policy
    to FSDP call
    
coremltools/test/optimize/torch/palettization/test_sensitive_k_means.py