def get_param_config(self, param_name: str, param: _torch.Tensor) -> KMeansConfig:
        """
        Returns KMeansConfig for the specified parameter
        """
        config = self.config[param_name]
        block_size = param.shape[config.axis] if config.block_size is None else config.block_size
        cluster_dim = 1 if config.cluster_dim is None else config.cluster_dim
        importance = self._reshape_for_kmeans(param_name, config.importance)
        mask = self._reshape_for_kmeans(param_name, config.mask)

        return KMeansConfig(
            n_bits=config.n_bits,
            axis=config.axis,
            lut_dtype=config.lut_dtype,
            block_size=block_size,
            cluster_dim=cluster_dim,
            enable_per_channel_scale=config.enable_per_channel_scale,
            mask=mask,
            importance=importance,
            enable_fast_kmeans_mode=config.enable_fast_kmeans_mode,
            rounding_precision=config.rounding_precision,
        )

        Returns KMeansConfig for the specified parameter
        
coremltools/optimize/torch/_utils/k_means.py