def _compute_lut_and_indices(self, param_name: str, param: _torch.Tensor):
        """
        Compute LUT and indices from parameter.
        For 4-bit palettization and param shape (32, 16, 3, 3),
        Case-1: If block_size = 4 and axis = 0, then LUT has shape (8, 1, 1, 1, 16, 1)
        Case-2: If block_size = 4 and axis = 1, then LUT has shape (1, 4, 1, 1, 16, 1)
        Case-3: If cluster_dim = 4, then LUT has shape (1, 1, 1, 1, 16, 4)
        """
        axis = self.config[param_name].axis
        num_channels = param.shape[axis]
        mask = self.config[param_name].mask
        block_size = self.config[param_name].block_size
        block_size = num_channels if block_size is None else block_size
        cluster_dim = self.config[param_name].cluster_dim
        orig_param_shape = self._parameter_metadata[param_name]["shape"]
        cluster_dim = 1 if cluster_dim is None else cluster_dim

        lut, indices = [], []
        if cluster_dim == 1:
            # Scalar palettization
            for block_idx in range(0, num_channels, block_size):
                if axis == 0:
                    lut_idx, ind_idx = _torch.unique(
                        param[block_idx : block_idx + block_size, :],
                        return_inverse=True,
                    )
                else:
                    lut_idx, ind_idx = _torch.unique(
                        param[:, block_idx : block_idx + block_size],
                        return_inverse=True,
                    )

                # Ensure param was correctly palettized
                # Unless a mask was applied, number of unique values cannot exceed 2^nbits
                max_unique_val = 2 ** self.config[param_name].n_bits
                assert mask is not None or len(lut_idx) <= max_unique_val, (
                    f"Found more than expected unique values in {self.module} "
                    f"for {param_name}, expected <= {max_unique_val}, found = {len(lut_idx)}"
                )
                # Pad lut with zeros if fewer than 2^n_bit unique values are found
                if len(lut_idx) < max_unique_val:
                    padded_lut_idx = _torch.zeros(max_unique_val)
                    padded_lut_idx[: len(lut_idx)] = lut_idx
                    lut_idx = padded_lut_idx

                lut.append(lut_idx)
                indices.append(ind_idx)

            lut = _torch.stack(lut).unsqueeze(1 - axis).unsqueeze(-1)
            indices = _torch.cat(indices, dim=axis)
            indices = self._reshape_to_original(param_name, indices)
        else:
            # Vector palettization
            # Reshape param for 2D clustering
            if axis == 0:
                param_reshaped = param.transpose(0, 1).reshape(-1, cluster_dim)
            else:
                param_reshaped = param.reshape(-1, cluster_dim)
            lut, indices = _torch.unique(param_reshaped, dim=0, return_inverse=True)

            # Undo reshaping in indices done for 2D clustering
            if axis == 0:
                indices = indices.reshape(param.shape[0] // cluster_dim, param.shape[1])
            else:
                indices = indices.reshape(param.shape[0], param.shape[1] // cluster_dim)

        # Incorporate param dimensions in lut shape
        for i in range(len(orig_param_shape) - lut.dim() + 2):
            lut = lut.unsqueeze(-3)

        return lut, indices

        Compute LUT and indices from parameter.
        For 4-bit palettization and param shape (32, 16, 3, 3),
        Case-1: If block_size = 4 and axis = 0, then LUT has shape (8, 1, 1, 1, 16, 1)
        Case-2: If block_size = 4 and axis = 1, then LUT has shape (1, 4, 1, 1, 16, 1)
        Case-3: If cluster_dim = 4, then LUT has shape (1, 1, 1, 1, 16, 4)
        
coremltools/optimize/torch/_utils/k_means.py