def test_g_idx_tp(self):
        """
        Test g_idx with autoTP and weight is split by K. For lowp_mode != INT8.
        Since it's complicated to write a test case for deepspeed,
        here we do the sharding manually to simulate the behavior.
        """

        class Mod(nn.Module):
            def __init__(self, ic, oc):
                super(Mod, self).__init__()
                # set bias to false and we add it ourselves
                self.linear = torch.nn.Linear(ic, oc, bias=False)

            def forward(self, x):
                return self.linear(x)

        shape_list = [[1, 32, 32], [32, 128, 128], [128, 128, 128], [1024, 128, 128]]
        group_size_list = [4, 16]
        lowp_mode_list = [WoqLowpMode.NONE, WoqLowpMode.FP16, WoqLowpMode.BF16]
        has_bias_list = [False, True]
        cases = itertools.product(
            shape_list, group_size_list, lowp_mode_list, has_bias_list
        )
        for shape, group_size, lowp_mode, has_bias in cases:
            bs, ic, oc = shape
            n_groups = ic // group_size
            dtype = torch.float
            if lowp_mode == WoqLowpMode.BF16:
                dtype = torch.bfloat16
            elif lowp_mode == WoqLowpMode.FP16:
                dtype = torch.half
            int4_weight = torch.randint(0, 15, (oc, ic), dtype=torch.uint8)
            packed_weight = (
                int4_weight[:, 1::2]
                .bitwise_left_shift(4)
                .bitwise_or_(int4_weight[:, ::2])
            )
            scales = torch.rand((oc, n_groups), dtype=torch.half) * 0.1
            zeros = torch.randint(6, 9, (oc, n_groups), dtype=torch.uint8)
            packed_zeros = torch.zeros(
                (oc, (n_groups * 4 + 32 - 1) // 32), dtype=torch.int32
            )
            for i in range(n_groups):
                packed_zeros[:, i // 8] = packed_zeros[:, i // 8].bitwise_or_(
                    zeros[:, i].int().bitwise_left_shift(4 * (i % 8))
                )
            g_idx = torch.arange(0, n_groups).to(torch.int).repeat(group_size)
            x = torch.randn((bs, ic), dtype=torch.float)

            # woq path
            b = torch.randn(oc).to(dtype) if has_bias else None
            qconfig_mapping = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                weight_dtype=WoqWeightDtype.INT4,
                lowp_mode=lowp_mode,
                group_size=group_size,
            )
            m = Mod(ic=ic // 2, oc=oc)
            woq_m_0 = copy.deepcopy(m)
            woq_m_0.linear.qconfig = qconfig_mapping.global_qconfig
            packed_weight_0 = packed_weight[
                :, : packed_weight.shape[1] // 2
            ].contiguous()
            g_idx_0 = g_idx[: g_idx.shape[0] // 2]
            woq_m_0.linear = (
                ipex.nn.modules.WeightOnlyQuantizedLinear.from_float_and_qweight(
                    woq_m_0.linear,
                    packed_weight_0,
                    WoqWeightDtype.INT4,
                    scales,
                    packed_zeros,
                    None,  # bias
                    group_size=group_size,
                    g_idx=g_idx_0,
                )
            )
            woq_m_1 = copy.deepcopy(m)
            woq_m_1.linear.qconfig = qconfig_mapping.global_qconfig
            packed_weight_1 = packed_weight[
                :, packed_weight.shape[1] // 2 :
            ].contiguous()
            g_idx_1 = g_idx[g_idx.shape[0] // 2 :]
            woq_m_1.linear = WeightOnlyQuantizedLinear.from_float_and_qweight(
                woq_m_1.linear,
                packed_weight_1,
                WoqWeightDtype.INT4,
                scales,
                packed_zeros,
                None,  # bias
                group_size=group_size,
                g_idx=g_idx_1,
            )
            x_0 = x[:, : x.shape[1] // 2].contiguous()
            y_0 = woq_m_0(x_0.to(dtype))
            x_1 = x[:, x.shape[1] // 2 :].contiguous()
            y_1 = woq_m_1(x_1.to(dtype))
            y = y_0 + y_1 + b if has_bias else y_0 + y_1

            # ref path
            scales_expanded = scales.repeat(1, group_size)
            zeros_expanded = zeros.repeat(1, group_size)
            dqw = (int4_weight.float() - zeros_expanded) * scales_expanded
            dqw = dqw.to(dtype)
            x = x.to(dtype)
            y_ref = torch.nn.functional.linear(x, dqw, bias=b).to(dtype)

            # check results
            atol = 1e-4
            rtol = 1e-5
            if lowp_mode == WoqLowpMode.BF16:
                atol = 1e-1
                rtol = 1e-3
            elif lowp_mode == WoqLowpMode.FP16:
                atol = 1.5e-2
                rtol = 1e-3
            torch.testing.assert_close(y, y_ref, atol=atol, rtol=rtol)

        Test g_idx with autoTP and weight is split by K. For lowp_mode != INT8.
        Since it's complicated to write a test case for deepspeed,
        here we do the sharding manually to simulate the behavior.
        
tests/cpu/test_quantization_default_recipe.py