def _construct_palettization_op(
        self,
        weight: np.ndarray,
        compression_info: CompressionInfo,
        name: str,
        compressed_var: Optional[Var] = None,
    ) -> Var:
        """
        The weight is constructed by 2**nbits unique values in each group.

        When `palettization_scale` is provided, it means the weight has scales before got palettized.
        More specifically, the diagram is:

        lut(fp16) \
                    -> constexpr_lut_to_dense -> dense(fp16) -> constexpr_blockwise_shift_scale -> dense(fp16)
        indices  /

        If compressed_var is not None, it's the var constructed by a previous compression function,
        which means this is a joint compression. For example, if the compression_info.compression_type
        is [CompressionType.PRUNING, CompressionType.PALETTIZATION], the compressed_var is the var
        produced by the pruning.
        """
        if compression_info.lut is None:
            raise ValueError("Missing lut in compression info. Please register a buffer for lut.")

        lut = compression_info.lut.detach().numpy()
        if len(lut.shape) == len(weight.shape) + 1:
            # The last dim to indicate vector size is by default 1 for scalar palettization.
            lut = np.expand_dims(lut, axis=-1)
        if len(lut.shape) != len(weight.shape) + 2:
            raise ValueError(
                f"In {name}, The rank of lut is invalid. It should match weight's rank where"
                f"lut.rank == weight.rank + 2). Got lut.rank {len(lut.shape)} and weight.rank {len(weight.shape)}"
            )

        num_palettes = lut.shape[-2]
        nbits = int(math.ceil(math.log2(num_palettes)))
        if 2**nbits != num_palettes:
            # Padding lut to make it has 2**nbits dim size on -2 axis.
            padding_shape = list(lut.shape)
            padding_shape[-2] = 2**nbits - num_palettes
            lut = np.concatenate([lut, np.zeros(padding_shape, dtype=lut.dtype)], axis=-2)
            num_palettes = lut.shape[-2]

        if compression_info.palettization_scale is not None:
            # The weight has scales, which means the palettization is on the pre-scale data.
            scale = compression_info.palettization_scale.detach().numpy()
            # For conv/conv_transpose, the weight has rank=4, so we auto-expand scale and zero-point if
            # it only has two elements.
            if len(weight.shape) == 4 and len(scale.shape) == 2:
                scale = np.expand_dims(np.expand_dims(scale, axis=-1), axis=-1)
            if len(scale.shape) != len(weight.shape):
                raise ValueError(
                    f"In {name}, the scale should have the same rank as weight, but got "
                    f"{scale.shape} vs {weight.shape}."
                )
            weight = weight / scale

        vector_axis = compression_info.vector_axis
        if lut.shape[-1] > 1:
            if vector_axis is None:
                # The cto.torch uses 0 for vector axis.
                logger.warning(
                    "It's recommended to provide vector_axis for vector palettization. "
                    "Defaulting to axis zero."
                )
                vector_axis = 0
        indices = optimize_utils.find_indices_for_lut(weight, lut, vector_axis)

        if CompressionType.QUANTIZATION.value in compression_info.compression_type:
            # In joint palettization + quantization, the `lut` in the palettization op should be
            # quantized, so we calculate the quantized lut on-the-fly.
            tmp_quant_var = self._construct_quantization_op(
                lut, compression_info, name + "_tmp_quant"
            )
            lut = tmp_quant_var.op.data.val

        if compressed_var is None:
            result = frontend_utils._construct_constexpr_lut_op(indices, lut, vector_axis, name)
        else:
            # Specially handles joint compression, such as using sparse op if joint with pruning.
            if compressed_var.op.op_type == "constexpr_sparse_to_dense":
                mask, nonzero_data = mb.constexpr_lut_to_sparse(
                    indices_mask=compressed_var.op.mask,
                    indices_nonzero_data=indices[compressed_var.op.mask.val != 0].flatten(),
                    lut=lut,
                    vector_axis=vector_axis,
                    before_op=compressed_var.op,
                    name=compressed_var.op.name + "_palettized",
                )
                result = mb.constexpr_sparse_to_dense(
                    nonzero_data=nonzero_data, mask=mask, name=name
                )
            else:
                raise ValueError(
                    "Unsupported joint compression combination. The palettization can only be joint "
                    f"with pruning, but got {compressed_var.op.op_type}. Please check the value of "
                    "'compression_type' in your registered buffers."
                )

        if compression_info.palettization_scale is not None:
            if not is_current_opset_version_compatible_with(_target.iOS18):
                raise ValueError(
                    "The palettization with per-channel-scale is only supported in iOS18+. Please "
                    "set the minimum_deployment_target to iOS18 or later."
                )
            result = mb.constexpr_blockwise_shift_scale(
                data=result, scale=scale, offset=None, name=name
            )
        return result

        The weight is constructed by 2**nbits unique values in each group.

        When `palettization_scale` is provided, it means the weight has scales before got palettized.
        More specifically, the diagram is:

        lut(fp16)                     -> constexpr_lut_to_dense -> dense(fp16) -> constexpr_blockwise_shift_scale -> dense(fp16)
        indices  /

        If compressed_var is not None, it's the var constructed by a previous compression function,
        which means this is a joint compression. For example, if the compression_info.compression_type
        is [CompressionType.PRUNING, CompressionType.PALETTIZATION], the compressed_var is the var
        produced by the pruning.
        
coremltools/converters/mil/frontend/torch/converter.py