def test_joint_prune_palettize_weights(
        self, compute_unit, backend, nbits, channel_group_size, prune_first
    ):
        """Jointly prune and palettize the model, where non-sparse entries are palettized."""
        model, inputs, torch_input_values, coreml_input_values = get_test_model_and_data_complex()
        torchmodel = torch.jit.trace(model, torch_input_values)
        mlmodel = ct.convert(
            torchmodel,
            inputs=inputs,
            convert_to="mlprogram",
            minimum_deployment_target=backend.opset_version,
            compute_precision=ct.precision.FLOAT16
            if backend.precision == "fp16"
            else ct.precision.FLOAT32,
            compute_units=compute_unit,
        )

        prune_config = cto.coreml.OptimizationConfig(
            global_config=cto.coreml.OpMagnitudePrunerConfig(
                target_sparsity=0.2,
                weight_threshold=500,
            )
        )
        palettize_config = cto.coreml.OptimizationConfig(
            global_config=cto.coreml.OpPalettizerConfig(
                mode="uniform",
                nbits=nbits,
                granularity="per_grouped_channel",
                group_size=channel_group_size,
                weight_threshold=500,
            )
        )

        pytest_context_manager = nullcontext()
        if backend.opset_version < ct.target.iOS18:
            pytest_context_manager = pytest.raises(ValueError, match=self._ERR_MSG)

        with pytest_context_manager:
            if prune_first:
                mlmodel_pruned = cto.coreml.prune_weights(mlmodel, prune_config)
                mlmodel_joint_pruned_palettized = cto.coreml.palettize_weights(
                    mlmodel_pruned, palettize_config, joint_compression=True
                )
            else:
                mlmodel_palettized = cto.coreml.palettize_weights(mlmodel, palettize_config)
                mlmodel_joint_pruned_palettized = cto.coreml.prune_weights(
                    mlmodel_palettized, prune_config, joint_compression=True
                )

        if backend.opset_version < ct.target.iOS18:
            return

        # If run prune first, the all-zero const for lstm won't have nonzero-data, so it won't be
        # further quantized.
        lstm_weight_compression_ops = (
            ["constexpr_sparse_to_dense"]
            if prune_first
            else ["constexpr_lut_to_sparse", "constexpr_sparse_to_dense"]
        )
        expected_ops = (
            ["constexpr_lut_to_sparse", "constexpr_sparse_to_dense", "conv"] * 2
            + ["reshape"]
            + ["constexpr_lut_to_sparse", "constexpr_sparse_to_dense", "linear"] * 2
            + lstm_weight_compression_ops
            + ["constexpr_lut_to_sparse", "constexpr_sparse_to_dense"] * 2
            + ["lstm"]
        )
        prog = mlmodel_joint_pruned_palettized._mil_program
        assert get_op_types_in_program(prog) == expected_ops

        for linear_op in prog.find_ops(op_type="linear"):
            assert linear_op.weight.op.op_type == "constexpr_sparse_to_dense"
        for conv_op in prog.find_ops(op_type="conv"):
            assert conv_op.weight.op.op_type == "constexpr_sparse_to_dense"

        sparse_palettize_ops = prog.find_ops(op_type="constexpr_lut_to_sparse")
        assert len(sparse_palettize_ops) > 0
        for sparse_palettize_op in sparse_palettize_ops:
            assert sparse_palettize_op.indices_nonzero_data.dtype == types.string_to_builtin(
                f"uint{nbits}"
            )
            assert sparse_palettize_op.indices_mask.dtype == types.uint1
            assert sparse_palettize_op.lut.dtype == types.fp16
            # Both outputs of the sparse_palettize_op should be used by the child constexpr_sparse_to_dense op.
            assert (
                sparse_palettize_op.outputs[0].child_ops[0].op_type == "constexpr_sparse_to_dense"
            )
            assert (
                sparse_palettize_op.outputs[1].child_ops[0].op_type == "constexpr_sparse_to_dense"
            )
            # As both palettization and pruning is on the original weight, the shape of lut should
            # match the original weight's shape except on the output channel.
            weight_shape = sparse_palettize_op.outputs[1].child_ops[0].outputs[0].shape
            expected_lut_shape = [1] * len(weight_shape) + [2**nbits] + [1]
            if channel_group_size > 0:
                expected_lut_shape[0] = weight_shape[0] // channel_group_size
            assert sparse_palettize_op.lut.shape == tuple(expected_lut_shape)

        sparse_ops = prog.find_ops(op_type="constexpr_sparse_to_dense")
        assert len(sparse_ops) > 0
        for sparse_op in sparse_ops:
            assert sparse_op.mask.dtype == types.uint1
            assert sparse_op.nonzero_data.dtype == types.fp16

        if _macos_version() >= (15, 0):
            verify_model_outputs(
                mlmodel, mlmodel_joint_pruned_palettized, coreml_input_values, atol=5e-4
            )
Jointly prune and palettize the model, where non-sparse entries are palettized.
coremltools/test/optimize/coreml/test_post_training_quantization.py