def _get_compression_metadata(
        self, param_name: str, param: _torch.Tensor, lut: _torch.Tensor
    ) -> _CompressionMetadata:
        """
        Return compression metadata to be stored in the model for this parameter
        """
        metadata = _CompressionMetadata(param_name)
        compression_type = ["palettization"]

        # LUT
        # Incorporate param dimensions in lut shape
        orig_param_shape = self._parameter_metadata[param_name]["shape"]
        for i in range(len(orig_param_shape) - lut.dim() + 2):
            lut = lut.unsqueeze(-3)
        metadata.lut = lut

        # Per channel scale
        if self.config[param_name].enable_per_channel_scale:
            per_channel_scale = self._parameter_metadata[param_name]["per_channel_scale"]
            reshaped_param = self._reshape_to_original(param_name, param)
            for _ in range(reshaped_param.dim() - per_channel_scale.dim()):
                per_channel_scale = per_channel_scale.unsqueeze(-1)
            metadata.palettization_scale = per_channel_scale
        # LUT quantization
        if self.config[param_name].lut_dtype is not None:
            dtype = self.config[param_name].lut_dtype
            compression_type.append("quantization")
            metadata.quantization_n_bits = get_n_bits_from_dtype(dtype)
            scale = self._parameter_metadata[param_name]["lut_quantization_scale"]
            # match scale rank to lut rank
            for i in range(metadata.lut.dim() - scale.dim()):
                scale = scale.unsqueeze(-1)
            metadata.quantization_scale = scale
            zp = self._parameter_metadata[param_name]["lut_quantization_zp"]
            if zp is not None:
                # match zp rank to lut rank
                for i in range(metadata.lut.dim() - zp.dim()):
                    zp = zp.unsqueeze(-1)
                metadata.zero_point = zp
        # vector axis for cluster_dim > 1
        cluster_dim = self.config[param_name].cluster_dim
        if cluster_dim is not None and cluster_dim > 1:
            metadata.vector_axis = self.config[param_name].axis
        # Compression type
        metadata.compression_type = compression_type
        return metadata

        Return compression metadata to be stored in the model for this parameter
        
coremltools/optimize/torch/_utils/k_means.py