def _get_lut_pcs_weight(shape: Tuple[int, ...], nbits=4, scale_axis: int = 0):
        """Get a specific shape of weight produced by lut with per-channel-scale (pcs)."""
        num_palette = 2**nbits
        np_dtype = types.nptype_from_builtin(types.string_to_builtin(f"uint{nbits}"))
        indices = np.arange(np.prod(shape)).reshape(shape).astype(np_dtype)
        lut_shape = shape + (num_palette, 1)
        lut = np.arange(np.prod(lut_shape)).reshape(lut_shape).astype(np.float16)

        lut_op = mb.constexpr_lut_to_dense(indices=indices, lut=lut)
        scale_shape = [1] * len(shape)
        scale_shape[scale_axis] = shape[scale_axis]
        scale_shape = tuple(scale_shape)
        scale_val = np.arange(1, np.prod(scale_shape) + 1).reshape(scale_shape).astype(np.float16)
        return mb.constexpr_blockwise_shift_scale(
            data=lut_op,
            scale=scale_val,
        )
Get a specific shape of weight produced by lut with per-channel-scale (pcs).
coremltools/converters/mil/mil/passes/tests/test_quantization_passes.py