def test_selective_mixed_precision(algorithm, expected_layer_indices, input_model, tmp_path):
    """Test SelectiveMixedPrecision pass with different algorithms."""
    config = {"algorithm": algorithm}
    p = create_pass_from_dict(SelectiveMixedPrecision, config, disable_search=True)

    output_model = p.run(input_model, str(tmp_path))

    # Check that mixed_precision_info was added
    assert "mixed_precision_info" in output_model.model_attributes
    expected_mp_info = {
        "default": {"bits": PrecisionBits.BITS4},
        "overrides": {
            "lm_head": {"bits": 8},
        },
    }
    for idx in expected_layer_indices:
        expected_mp_info["overrides"][f"model.layers.{idx}.self_attn.v_proj"] = {"bits": PrecisionBits.BITS8}
        expected_mp_info["overrides"][f"model.layers.{idx}.mlp.down_proj"] = {"bits": PrecisionBits.BITS8}
    assert output_model.model_attributes["mixed_precision_info"] == expected_mp_info
Test SelectiveMixedPrecision pass with different algorithms.
test/unit_test/passes/pytorch/test_selective_mixed_precision.py