def test_sparse_gpt_metadata(config):
    """
    Test registration of metadata buffers for GPTQ algorithm
    """
    # Setup to get compressed model
    model = nn.Sequential(nn.Linear(4096, 1024))
    compressor_config = LayerwiseCompressorConfig.from_dict(config)
    compressor = LayerwiseCompressor(model, compressor_config)

    def calibration_loader():
        yield torch.rand(1, 4096)

    compressed_model = compressor.compress(calibration_loader(), device="cpu")

    # Extract registered metadata from state_dict
    state_dict = compressed_model[0].state_dict()
    metadata_dict = CompressionMetadata.from_state_dict(state_dict)
    assert len(metadata_dict) == 1
    assert "weight" in metadata_dict

    # Verification
    metadata = metadata_dict["weight"]
    if compressor_config.global_config.enable_normal_float:
        assert metadata.compression_type == [
            CompressionType.pruning.value,
            CompressionType.palettization.value,
        ]
        assert metadata.lut.shape == (
            1,
            1,
            2**compressor_config.global_config.weight_n_bits,
            1,
        )
        assert metadata.palettization_scale.shape == (state_dict["weight"].shape[0], 1)
    elif (
        compressor_config.global_config.weight_n_bits is not None
        and compressor_config.global_config.weight_n_bits < 16
    ):
        assert metadata.compression_type == [
            CompressionType.pruning.value,
            CompressionType.quantization.value,
        ]
        assert metadata.quantization_n_bits == compressor_config.global_config.weight_n_bits
        assert metadata.zero_point.shape == metadata.quantization_scale.shape

    assert METADATA_VERSION_BUFFER in compressed_model.state_dict()
    assert torch.equal(compressed_model.state_dict()[METADATA_VERSION_BUFFER], METADATA_VERSION)

    Test registration of metadata buffers for GPTQ algorithm
    
coremltools/test/optimize/torch/layerwise_compression/test_algorithms.py