def _weight_int4pack_mm(context, node):
    """
    The first argument is the same as torch.mm, but the second argument (weight) is packed.
    The packed weight has rank=4, because the meta registration in dynamo requires operator has the same output shape
    for each device. So it creates a fake shape {N / 8, K / (16 * innerKTiles), 32, innerKTiles / 2} for CPU.

    More specifically:

        # Original torch.mm
        torch.mm(a, b)

        # The int4 packed version mm
        b_uint8, b_scales_and_zeros = _group_quantize_tensor(
            b, n_bit=4, q_group_size=q_group
        )
        b_int4pack = torch._convert_weight_to_int4pack(
            b_uint8, inner_k_tiles
        )
        weight_int4pack_mm(a, b_int4pack, b_scales_and_zeros)
    """
    if Version(_torch.__version__) < Version("2.4.0"):
        raise AssertionError("To lower _weight_int4pack_mm, requires torch >= 2.4.0")

    logger.warning(
        "The current conversion of `_weight_int4pack_mm` op only works with model produced by torchao. "
        "If the op is produced by other libs, you may observe large numerical discrepancy."
    )

    if not _HAS_TORCHAO:
        raise AssertionError(
            f"{MSG_TORCHAO_NOT_FOUND}\n torchao is needed to convert torch blockwise quantized model."
        )

    inputs = _get_inputs(context, node, expected=4)
    x = inputs[0]
    y_int4pack = inputs[1].val
    group_size = inputs[2].val
    y_scales_and_zeros = inputs[3].val

    if y_int4pack is None or group_size is None or y_scales_and_zeros is None:
        raise NotImplementedError(
            "For `_weight_int4pack_mm` op, we only support static case, where y_int4pack, "
            "group_size, y_scales_and_zeros are all known during compilation time."
        )

    if not (len(y_scales_and_zeros.shape) == 3 and y_scales_and_zeros.shape[2] == 2):
        raise ValueError(
            "The scales_and_zeros from torch should have 3 dims and last dim has size 2."
        )
    scales = _np.transpose(y_scales_and_zeros[:, :, 0])
    zero_points = _np.transpose(y_scales_and_zeros[:, :, 1])

    if _np.allclose(zero_points, zero_points.astype("int32")):
        zero_points = zero_points.astype("int32")
    else:
        zero_points = zero_points.astype(_np.float32)

    # Unpack the result of `torch._convert_weight_to_int4pack` back to plain layout.
    # TODO: Use `torchao.ops.unpack_tensor_core_tiled_layout` to unpack after it has CPU implementation.
    # The current way to unpack by using _weight_int4pack_mm with eye matrix is a workaround on CPU.
    if len(y_int4pack.shape) != 4:
        raise ValueError(
            f"The packed y from torch should have 4 dims, but got {len(y_int4pack.shape)}."
        )
    inner_k_tiles = y_int4pack.shape[-1] * 2
    y_unpacked_shape = (y_int4pack.shape[0] * 8, y_int4pack.shape[1] * (inner_k_tiles * 16))
    eye_shape = y_unpacked_shape[1]
    quant_min = 0
    quant_max = 2**4 - 1
    with _torch.no_grad():
        y_dequantized = (
            _torch._weight_int4pack_mm(
                _torch.eye(eye_shape, device=_torch.device("cpu"), dtype=_torch.float32),
                _torch.from_numpy(y_int4pack),
                group_size,
                _torch.from_numpy(y_scales_and_zeros).float(),
            )
            .t()
            .contiguous()
        )
        zero_point_domain = (
            torchao_quant.ZeroPointDomain.INT
            if _np.issubdtype(zero_points.dtype, _np.integer)
            else torchao_quant.ZeroPointDomain.FLOAT
        )
        y_quantized = torchao_quant.quantize_affine(
            y_dequantized,
            (1, group_size),
            _torch.from_numpy(scales),
            _torch.from_numpy(zero_points),
            _torch.int32,
            quant_min=quant_min,
            quant_max=quant_max,
            zero_point_domain=zero_point_domain,
        )
    y_quantized = y_quantized.numpy().astype(_np.uint8)
    if len(y_quantized.shape) != 2:
        raise ValueError(
            f"The unpacked quantized y should have 2 dims, but got {len(y_quantized.shape)}."
        )
    if not _np.logical_and(y_quantized >= 0, y_quantized <= 15).all():
        raise ValueError("All elements should be within 4-bit range (0 to 15).")

    # If zero_point_domain in `quantize_affine` is set to `ZeroPointDomain.INT`, it matches with CoreML implementation:
    #         quant = torch.clamp(torch.round(input * (1.0 / scale)) + zero_point, quant_min, quant_max)
    # However, for `ZeroPointDomain.FLOAT`, torchao did following transformations to make it compatible with `tinygemm`:
    #         mid_point = (quant_max + quant_min + 1) / 2
    #         min_val = zero_point - scale * mid_point
    #         quant = torch.clamp(torch.round((input - min_val) / scale), quant_min, quant_max))
    # As we want to make sure the quantize matches CoreML dequant op, we have to do following transformations:
    #         dequant = (quant - mid_point) * scale + zp
    # so we can re-write the expression as
    #         dequant = (quant - (mid_point - zp / scale)) * scale
    # which means the zero_point in CoreML is actually `mid_point - zp / scale`.
    if not _np.issubdtype(zero_points.dtype, _np.integer):
        mid_point = (quant_max + quant_min + 1) / 2
        zero_points = mid_point - zero_points / scales

    # Use MIL constexpr op to represent the quantization.
    quantized_np_dtype = types.nptype_from_builtin(types.string_to_builtin("uint4"))
    dequant_weights = _utils._construct_constexpr_dequant_op(
        y_quantized.astype(quantized_np_dtype), zero_points, scales, axis=-1, name=inputs[1].name
    )

    res = mb.linear(x=x, weight=dequant_weights, name=node.name)
    context.add(res)

    The first argument is the same as torch.mm, but the second argument (weight) is packed.
    The packed weight has rank=4, because the meta registration in dynamo requires operator has the same output shape
    for each device. So it creates a fake shape {N / 8, K / (16 * innerKTiles), 32, innerKTiles / 2} for CPU.

    More specifically:

        # Original torch.mm
        torch.mm(a, b)

        # The int4 packed version mm
        b_uint8, b_scales_and_zeros = _group_quantize_tensor(
            b, n_bit=4, q_group_size=q_group
        )
        b_int4pack = torch._convert_weight_to_int4pack(
            b_uint8, inner_k_tiles
        )
        weight_int4pack_mm(a, b_int4pack, b_scales_and_zeros)
    
coremltools/converters/mil/frontend/torch/quantization_ops.py