def test_unpack_int4packed_by_mm_with_eye_matrix(self, use_numpy, inner_k_tiles, group_size):
        """
        Check if the packed weight could be restored by _weight_int4pack_mm with eye matrix on CPU.

        As there is no kernel implemented for CPU to unpack the data packed by `torch._convert_weight_to_int4pack`,
        we use `torch._weight_int4pack_mm` to do matrix multiplication with an eye matrix to get unpacked data.
        """
        if use_numpy:
            y_np = numpy.random.rand(128, 128).astype(np.float32)
            y = torch.from_numpy(y_np).to(torch.device("cpu"))
        else:
            y = torch.rand(128, 128, dtype=torch.float32, device=torch.device("cpu"))

        (
            y_quantized,
            y_scales_and_zeros,
        ) = torchao.quantization.utils.groupwise_affine_quantize_tensor(
            y, n_bit=4, groupsize=group_size, dtype=torch.float32
        )
        y_int4packed = torch._convert_weight_to_int4pack(y_quantized, inner_k_tiles)
        y_unpacked_shape = (y_int4packed.shape[0] * 8, y_int4packed.shape[1] * (inner_k_tiles * 16))
        eye_shape = y_unpacked_shape[1]
        eye_matrix = torch.eye(eye_shape, device=torch.device("cpu"), dtype=torch.float32)
        if Version(torch.__version__) < Version("2.4.0"):
            # The `torch._weight_int4pack_mm` op requires bfloat16 before PyTorch 2.4.0.
            eye_matrix = eye_matrix.to(torch.bfloat16)
            y_scales_and_zeros = y_scales_and_zeros.to(torch.bfloat16)
        y_dequant = torch._weight_int4pack_mm(
            eye_matrix,
            y_int4packed,
            group_size,
            y_scales_and_zeros,
        )
        y_dequant = y_dequant.t().contiguous().float()

        # Makes sure this `_weight_int4pack_mm` with eye matrix fully restores the original y.
        np.testing.assert_allclose(y_dequant.numpy(), y.numpy(), atol=0.035, rtol=0.05)

        # Also verifies that the quantized y could be accurately reproduced by torchao utils.
        scales = torch.transpose(y_scales_and_zeros[:, :, 0], 0, 1)
        zero_points = torch.transpose(y_scales_and_zeros[:, :, 1], 0, 1)
        block_size = (1, group_size)
        y_dequant_quantized = torchao_quant.quantize_affine(
            y_dequant,
            block_size,
            scales,
            zero_points,
            torch.int32,
            quant_min=0,
            quant_max=2**4 - 1,
            zero_point_domain=torchao_quant.ZeroPointDomain.FLOAT,
        )
        assert torch.equal(y_quantized, y_dequant_quantized)

        # The torchao dequantization utils should be able to recover the original y.
        y_dequantized_by_torchao = torchao_quant.dequantize_affine(
            y_quantized,
            (1, group_size),
            scales,
            zero_points,
            torch.int32,
            quant_min=0,
            quant_max=2**4 - 1,
            zero_point_domain=torchao_quant.ZeroPointDomain.FLOAT,
        )
        np.testing.assert_allclose(y_dequant.numpy(), y_dequantized_by_torchao.numpy(), rtol=4e-3)

        Check if the packed weight could be restored by _weight_int4pack_mm with eye matrix on CPU.

        As there is no kernel implemented for CPU to unpack the data packed by `torch._convert_weight_to_int4pack`,
        we use `torch._weight_int4pack_mm` to do matrix multiplication with an eye matrix to get unpacked data.
        
coremltools/converters/mil/frontend/torch/test/test_torch_quantization_ops.py