jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_prepare_prompt():
    model_runner = ModelRunner(None, None, None)
    model_runner.set_block_size(16)

    batch_size = random.randint(1, 256)
    prompt_lens = []
    seq_group_metadata_list = []
    for i in range(batch_size):
        # make sure all tokens fit into one block
        prompt_len = i % (model_runner.block_size - 1) + 1
        prompt_lens.append(prompt_len)
        seq_data = list(range(prompt_len))
        seq_group_metadata_list.append(
            SequenceGroupMetadata(
                request_id=f"test_{i}",
                is_prompt=True,
                seq_data={0: SequenceData(seq_data)},
                sampling_params=SamplingParams(temperature=0),
                block_tables={0: [1]},
            ))

    expected_selected_token_indices = []
    selected_token_start_idx = 0
    max_seq_len = max(prompt_lens)
    for prompt_len in prompt_lens:
        expected_selected_token_indices.append(selected_token_start_idx +
                                               prompt_len - 1)
        selected_token_start_idx += max_seq_len
    input_tokens, input_positions, _ = model_runner._prepare_prompt(
        seq_group_metadata_list)
    sampling_metadata = model_runner._prepare_sample(seq_group_metadata_list,
                                                     prompt_lens)
    assert input_tokens.shape == (batch_size, max_seq_len)
    assert input_positions.shape == (batch_size, max_seq_len)
    torch.testing.assert_close(input_tokens, input_positions)

    actual = sampling_metadata.selected_token_indices
    expected = torch.tensor(expected_selected_token_indices,
                            device=actual.device,
                            dtype=actual.dtype)
    torch.testing.assert_close(actual, expected)
//  [python]
def test_add_cnn_layer_rainbow(device):
    evolvable_cnn = EvolvableCNN(
        input_shape=[1, 128, 128],
        channel_size=[32, 32],
        kernel_size=[3, 3],
        stride_size=[1, 1],
        hidden_size=[64, 64],
        num_actions=10,
        rainbow=True,
        layer_norm=True,
        support=torch.linspace(0.0, 200.0, 51).to(device),
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_add_cnn_layer_rainbow(device):
    evolvable_cnn = EvolvableCNN(
        input_shape=[1, 128, 128],
        channel_size=[32, 32],
        kernel_size=[3, 3],
        stride_size=[1, 1],
        hidden_size=[64, 64],
        num_actions=10,
        rainbow=True,
        layer_norm=True,
        support=torch.linspace(0.0, 200.0, 51).to(device),
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.add_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_remove_cnn_layer_rainbow(device):
    evolvable_cnn = EvolvableCNN(
        input_shape=[1, 16, 16],
        channel_size=[32, 32],
        kernel_size=[3, 3],
        stride_size=[1, 1],
        hidden_size=[64, 64],
        num_actions=10,
        rainbow=True,
        layer_norm=True,
        support=torch.linspace(0.0, 200.0, 51).to(device),
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num - 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_remove_cnn_layer_rainbow(device):
    evolvable_cnn = EvolvableCNN(
        input_shape=[1, 16, 16],
        channel_size=[32, 32],
        kernel_size=[3, 3],
        stride_size=[1, 1],
        hidden_size=[64, 64],
        num_actions=10,
        rainbow=True,
        layer_norm=True,
        support=torch.linspace(0.0, 200.0, 51).to(device),
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num < 6:
        assert len(evolvable_cnn.channel_size) == initial_channel_num - 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_remove_cnn_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num > 1:
        assert len(evolvable_cnn.channel_size) == initial_channel_num - 1
        for key, param in new_net.named_parameters():
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_remove_cnn_layer(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    initial_channel_num = len(evolvable_cnn.channel_size)
    initial_net = evolvable_cnn.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_cnn.remove_cnn_layer()
    new_net = evolvable_cnn.feature_net
    if initial_channel_num > 1:
        assert len(evolvable_cnn.channel_size) == initial_channel_num - 1
        for key, param in new_net.named_parameters():
            if (
                key in initial_net_dict.keys()
                and param.shape == initial_net_dict[key].shape
            ):
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_cnn.channel_size) == initial_channel_num
//  [python]
def test_clone_instance(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    original_feature_net_dict = dict(evolvable_cnn.feature_net.named_parameters())
    original_value_net_dict = dict(evolvable_cnn.value_net.named_parameters())
    clone = evolvable_cnn.clone()
    clone_feature_net = clone.feature_net
    clone_value_net = clone.value_net
    assert isinstance(clone, EvolvableCNN)
    assert clone.init_dict == evolvable_cnn.init_dict
    assert str(clone.state_dict()) == str(evolvable_cnn.state_dict())
    for key, param in clone_feature_net.named_parameters():
        torch.testing.assert_allclose(param, original_feature_net_dict[key])
    for key, param in clone_value_net.named_parameters():
        torch.testing.assert_allclose(param, original_value_net_dict[key])
//  [python]
def test_clone_instance(
    input_shape,
    channel_size,
    kernel_size,
    stride_size,
    hidden_size,
    num_actions,
    device,
):
    evolvable_cnn = EvolvableCNN(
        input_shape=input_shape,
        channel_size=channel_size,
        kernel_size=kernel_size,
        stride_size=stride_size,
        hidden_size=hidden_size,
        num_actions=num_actions,
        device=device,
    )
    original_feature_net_dict = dict(evolvable_cnn.feature_net.named_parameters())
    original_value_net_dict = dict(evolvable_cnn.value_net.named_parameters())
    clone = evolvable_cnn.clone()
    clone_feature_net = clone.feature_net
    clone_value_net = clone.value_net
    assert isinstance(clone, EvolvableCNN)
    assert clone.init_dict == evolvable_cnn.init_dict
    assert str(clone.state_dict()) == str(evolvable_cnn.state_dict())
    for key, param in clone_feature_net.named_parameters():
        torch.testing.assert_close(param, original_feature_net_dict[key])
    for key, param in clone_value_net.named_parameters():
        torch.testing.assert_close(param, original_value_net_dict[key])
//  [python]
def test_weight_only_quantization_new_gelu_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)
                self.gelu = nn.GELU(approximate="tanh")

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

        bias_list = [False, True]
        bf16_list = [False, True]
        cases = itertools.product(bias_list, bf16_list)
        for bias, bf16 in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data)
                    output2 = torch.ops.torch_ipex.woq_linear_new_gelu(
                        data, woq_model.linear._op_context.get_data_handle()
                    )
                    torch.testing.assert_allclose(
                        output1, output2.to(output1.dtype), atol=1e-2, rtol=1e-4
                    )
//  [python]
def test_weight_only_quantization_new_gelu_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)
                self.gelu = nn.GELU(approximate="tanh")

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

        bias_list = [False, True]
        bf16_list = [False, True]
        cases = itertools.product(bias_list, bf16_list)
        for bias, bf16 in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data)
                    output2 = torch.ops.torch_ipex.woq_linear_new_gelu(
                        data, woq_model.linear._op_context.get_data_handle()
                    )
                    torch.testing.assert_close(
                        output1, output2.to(output1.dtype), atol=1e-2, rtol=1e-4
                    )
//  [python]
def test_weight_only_quantization_gelu_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)
                self.gelu = nn.GELU()

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

        bias_list = [False, True]
        bf16_list = [False, True]
        cases = itertools.product(bias_list, bf16_list)
        for bias, bf16 in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data)
                    output2 = torch.ops.torch_ipex.woq_linear_gelu(
                        data, woq_model.linear._op_context.get_data_handle()
                    )
                    torch.testing.assert_allclose(
                        output1, output2.to(output1.dtype), atol=1e-2, rtol=1e-4
                    )
//  [python]
def test_weight_only_quantization_gelu_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)
                self.gelu = nn.GELU()

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

        bias_list = [False, True]
        bf16_list = [False, True]
        cases = itertools.product(bias_list, bf16_list)
        for bias, bf16 in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data)
                    output2 = torch.ops.torch_ipex.woq_linear_gelu(
                        data, woq_model.linear._op_context.get_data_handle()
                    )
                    torch.testing.assert_close(
                        output1, output2.to(output1.dtype), atol=1e-2, rtol=1e-4
                    )
//  [python]
def test_gru(self):
        input_dim = 32
        gru = GRUGatingUnit(input_dim, 1.)
        x = torch.rand((4, 12, 32))
        y = torch.rand((4, 12, 32))
        out = gru(x, y)
        assert out.shape == x.shape
        gru = GRUGatingUnit(input_dim, 100000.)  # set high bias to check 'out' is similar to the first input 'x'
        # In GTrXL the bias is initialized with a value high enough such that information coming from the second input
        # 'y' are partially ignored so to produce a behavior more similar to a MDP, thus giving less importance to
        # past information
        out = gru(x, y)
        torch.testing.assert_allclose(out, x)
//  [python]
def test_gru(self):
        input_dim = 32
        gru = GRUGatingUnit(input_dim, 1.)
        x = torch.rand((4, 12, 32))
        y = torch.rand((4, 12, 32))
        out = gru(x, y)
        assert out.shape == x.shape
        gru = GRUGatingUnit(input_dim, 100000.)  # set high bias to check 'out' is similar to the first input 'x'
        # In GTrXL the bias is initialized with a value high enough such that information coming from the second input
        # 'y' are partially ignored so to produce a behavior more similar to a MDP, thus giving less importance to
        # past information
        out = gru(x, y)
        torch.testing.assert_close(out, x)
//  [python]
def test_weight_only_quantization_jit_save_load(self):
        class M(nn.Module):
            def __init__(self, input_channel, output_channel, has_bias):
                super(M, self).__init__()
                self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)

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

        def test(feature, has_bias, w_dtype):
            model = M(feature[1], feature[2], has_bias)
            m = model.eval()
            example_inputs = torch.rand(feature[0], feature[1])

            qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                weight_dtype=w_dtype
            )
            prepared_model = prepare(
                m, qconfig, example_inputs=example_inputs, inplace=False
            )
            with torch.no_grad():
                converted_model = convert(prepared_model)

                with tempfile.NamedTemporaryFile() as fp:
                    # save
                    with torch.no_grad():
                        traced_model = torch.jit.trace(converted_model, example_inputs)
                        traced_model = torch.jit.freeze(traced_model)
                        traced_model.save(fp.name)

                    # load
                    loaded_model = torch.jit.load(fp.name)

                    # Compare results of original model and loaded model
                    output_ref = traced_model(example_inputs)
                    output = loaded_model(example_inputs)
                    torch.testing.assert_allclose(output_ref, output)

        shape_list = [
            [3, 31, 31],
            [4, 4096, 4096],
            [9, 4095, 4095],
            [196, 4095, 16383],
        ]
        use_bias_list = [True, False]
        w_dtype_list = [WoqWeightDtype.INT8, WoqWeightDtype.INT4]
        cases = itertools.product(shape_list, use_bias_list, w_dtype_list)
        for shape, use_bias, w_dtype in cases:
            test(shape, use_bias, w_dtype)
//  [python]
def test_weight_only_quantization_jit_save_load(self):
        class M(nn.Module):
            def __init__(self, input_channel, output_channel, has_bias):
                super(M, self).__init__()
                self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)

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

        def test(feature, has_bias, w_dtype):
            model = M(feature[1], feature[2], has_bias)
            m = model.eval()
            example_inputs = torch.rand(feature[0], feature[1])

            qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                weight_dtype=w_dtype
            )
            prepared_model = prepare(
                m, qconfig, example_inputs=example_inputs, inplace=False
            )
            with torch.no_grad():
                converted_model = convert(prepared_model)

                with tempfile.NamedTemporaryFile() as fp:
                    # save
                    with torch.no_grad():
                        traced_model = torch.jit.trace(converted_model, example_inputs)
                        traced_model = torch.jit.freeze(traced_model)
                        traced_model.save(fp.name)

                    # load
                    loaded_model = torch.jit.load(fp.name)

                    # Compare results of original model and loaded model
                    output_ref = traced_model(example_inputs)
                    output = loaded_model(example_inputs)
                    torch.testing.assert_close(output_ref, output)

        shape_list = [
            [3, 31, 31],
            [4, 4096, 4096],
            [9, 4095, 4095],
            [196, 4095, 16383],
        ]
        use_bias_list = [True, False]
        w_dtype_list = [WoqWeightDtype.INT8, WoqWeightDtype.INT4]
        cases = itertools.product(shape_list, use_bias_list, w_dtype_list)
        for shape, use_bias, w_dtype in cases:
            test(shape, use_bias, w_dtype)
//  [python]
def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            self.assertClose(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)
//  [python]
def assertClose(self, cutensor_tensor, torch_tensor):
        self.assertEqual(cutensor_tensor.shape, torch_tensor.shape)
        self.assertEqual(torch.is_complex(cutensor_tensor), torch.is_complex(torch_tensor))
        if torch.is_complex(cutensor_tensor):
            self.assertClose(torch.real(cutensor_tensor), torch.real(torch_tensor))
            self.assertClose(torch.imag(cutensor_tensor), torch.imag(torch_tensor))
        else:
            torch.testing.assert_close(cutensor_tensor, torch_tensor, rtol=5e-3, atol=6e-3)
//  [python]
def test_ddp_comm_hook_allreduce_hook(self):
        """
        This unit test verifies the ``allreduce`` hook registered case gives same result
        with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.ALLREDUCE)

        self.assertEqual(reference_grads, hook_grads)
//  [python]
def test_ddp_comm_hook_allreduce_hook(self):
        """
        This unit test verifies the ``allreduce`` hook registered case gives same result
        with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.ALLREDUCE)

        torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=0)
//  [python]
def test_weight_only_quantization_add_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)

            def forward(self, x, others):
                y = self.linear(x)
                for o in others:
                    y = torch.add(y, o)
                return y

        bias_list = [False, True]
        bf16_list = [False, True]
        others_len_list = [1, 2]
        cases = itertools.product(bias_list, bf16_list, others_len_list)
        for bias, bf16, others_len in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                others = [torch.rand(4, 64)] * others_len
                fused_op = (
                    torch.ops.torch_ipex.woq_linear_add
                    if others_len == 1
                    else torch.ops.torch_ipex.woq_linear_add_add
                )
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data, others)
                    output2 = fused_op(
                        data, woq_model.linear._op_context.get_data_handle(), others
                    )
                    torch.testing.assert_allclose(
                        output1, output2.to(output1.dtype), atol=1.5e-2, rtol=1e-3
                    )
//  [python]
def test_weight_only_quantization_add_fused_op(self):
        class Mod(nn.Module):
            def __init__(self, bias):
                super().__init__()
                self.linear = nn.Linear(64, 64, bias=bias)

            def forward(self, x, others):
                y = self.linear(x)
                for o in others:
                    y = torch.add(y, o)
                return y

        bias_list = [False, True]
        bf16_list = [False, True]
        others_len_list = [1, 2]
        cases = itertools.product(bias_list, bf16_list, others_len_list)
        for bias, bf16, others_len in cases:
            with torch.cpu.amp.autocast(
                enabled=bf16, dtype=torch.bfloat16 if bf16 else None
            ):
                model = Mod(bias).eval()
                data = torch.rand(4, 64)
                others = [torch.rand(4, 64)] * others_len
                fused_op = (
                    torch.ops.torch_ipex.woq_linear_add
                    if others_len == 1
                    else torch.ops.torch_ipex.woq_linear_add_add
                )
                qconfig = ipex.quantization.get_weight_only_quant_qconfig_mapping(
                    lowp_mode=2
                )
                prepared_model = prepare(
                    model, qconfig, example_inputs=data, inplace=False
                )
                with torch.no_grad():
                    woq_model = convert(prepared_model)
                    output1 = woq_model(data, others)
                    output2 = fused_op(
                        data, woq_model.linear._op_context.get_data_handle(), others
                    )
                    torch.testing.assert_close(
                        output1, output2.to(output1.dtype), atol=1.5e-2, rtol=1e-3
                    )
//  [python]
def test_ddp_comm_hook_fp16compress_hook(self):
        """
        This unit test verifies the ``fp16 compress`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.FP16_COMPRESS)

        self.assertEqual(reference_grads, hook_grads)
//  [python]
def test_ddp_comm_hook_fp16compress_hook(self):
        """
        This unit test verifies the ``fp16 compress`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.FP16_COMPRESS)

        torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=1e-4)
//  [python]
def test_ddp_comm_hook_quantize_per_tensor_hook(self):
        """
        This unit test verifies the ``quantize per tensor`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.QUANTIZE_PER_TENSOR)

        self.assertEqual(hook_grads, reference_grads)
//  [python]
def test_ddp_comm_hook_quantize_per_tensor_hook(self):
        """
        This unit test verifies the ``quantize per tensor`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.QUANTIZE_PER_TENSOR)

        torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=1e-4)
//  [python]
def test_ddp_comm_hook_quantize_per_channel_hook(self):
        """
        This unit test verifies the ``quantize per channel`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(
            process_group, DDPCommHookType.QUANTIZE_PER_CHANNEL
        )

        t.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=1e-4)
//  [python]
def test_ddp_comm_hook_quantize_per_channel_hook(self):
        """
        This unit test verifies the ``quantize per channel`` hook registered case
        gives close result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(
            process_group, DDPCommHookType.QUANTIZE_PER_CHANNEL
        )

        torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=1e-4)
//  [python]
def test_ddp_comm_hook_noop_hook(self):
        """
        This unit test verifies the ``noop`` hook registered case and a subsequent allreduce
        gives same result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.NOOP)
        # Apply a subsequent allreduce to average grads.
        hook_grads.div_(self.world_size)
        dist.all_reduce(hook_grads, group=process_group)

        self.assertEqual(hook_grads, reference_grads)
//  [python]
def test_ddp_comm_hook_noop_hook(self):
        """
        This unit test verifies the ``noop`` hook registered case and a subsequent allreduce
        gives same result with no hook registered case.
        """
        process_group = self._get_process_group_nccl()

        # No hook registered case, get the reference grads.
        reference_grads = self._get_grads(process_group, None)
        # Register hook case, get the hook grads.
        hook_grads = self._get_grads(process_group, DDPCommHookType.NOOP)
        # Apply a subsequent allreduce to average grads.
        hook_grads.div_(self.world_size)
        dist.all_reduce(hook_grads, group=process_group)

        torch.testing.assert_close(hook_grads, reference_grads, rtol=1e-5, atol=0)
//  [python]
def test_activation_checkpointing():
    dtype = torch.float16
    device = torch.device("cuda")
    test_model = DummyModel(is_checkpointed=True)
    ref_model = DummyModel(is_checkpointed=False)
    for model in [test_model, ref_model]:
        model.to(device=device, dtype=dtype)

    # copy weights
    test_model.load_state_dict(ref_model.state_dict())
    assert test_model.checkpointed_model.is_checkpointed is True
    assert ref_model.checkpointed_model.is_checkpointed is False

    # generate random input
    x = torch.randn(10, 10, device=device, dtype=dtype)

    # Forward pass
    with torch.random.fork_rng(devices=["cuda"]):
        ref_output = ref_model(x)
    checkpointed_output = test_model(x)
    assert test_model.checkpointed_model.fwd_counter == 1
    torch.testing.assert_allclose(checkpointed_output, ref_output)

    # Backward pass (check that fwd is called twice, and that we don't store the activations)
    ref_output.sum().backward()
    assert ref_model.checkpointed_model.fwd_counter == 1, "ref_model fwd should not be called twice"

    # make sure grads are not synced between test_model and ref_model
    assert ref_model.dense0.weight.grad is not None
    assert test_model.dense0.weight.grad is None

    assert test_model.checkpointed_model.fwd_counter == 1
    checkpointed_output.sum().backward()
    assert test_model.checkpointed_model.fwd_counter == 2, "test_model fwd should be called twice"

    # compare all models grads
    for ref_param, checkpointed_param in zip(ref_model.parameters(), test_model.parameters()):
        torch.testing.assert_allclose(ref_param.grad, checkpointed_param.grad)
//  [python]
def test_activation_checkpointing():
    dtype = torch.float16
    device = torch.device("cuda")
    test_model = DummyModel(is_checkpointed=True)
    ref_model = DummyModel(is_checkpointed=False)
    for model in [test_model, ref_model]:
        model.to(device=device, dtype=dtype)

    # copy weights
    test_model.load_state_dict(ref_model.state_dict())
    assert test_model.checkpointed_model.is_checkpointed is True
    assert ref_model.checkpointed_model.is_checkpointed is False

    # generate random input
    x = torch.randn(10, 10, device=device, dtype=dtype)

    # Forward pass
    with torch.random.fork_rng(devices=["cuda"]):
        ref_output = ref_model(x)
    checkpointed_output = test_model(x)
    assert test_model.checkpointed_model.fwd_counter == 1
    torch.testing.assert_close(checkpointed_output, ref_output)

    # Backward pass (check that fwd is called twice, and that we don't store the activations)
    ref_output.sum().backward()
    assert ref_model.checkpointed_model.fwd_counter == 1, "ref_model fwd should not be called twice"

    # make sure grads are not synced between test_model and ref_model
    assert ref_model.dense0.weight.grad is not None
    assert test_model.dense0.weight.grad is None

    assert test_model.checkpointed_model.fwd_counter == 1
    checkpointed_output.sum().backward()
    assert test_model.checkpointed_model.fwd_counter == 2, "test_model fwd should be called twice"

    # compare all models grads
    for ref_param, checkpointed_param in zip(ref_model.parameters(), test_model.parameters()):
        torch.testing.assert_close(ref_param.grad, checkpointed_param.grad)
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
      assert quant_type == torch.int8 or quant_type == torch.quint4x2

      lower_bound = -128 if quant_type == torch.int8 else -8
      upper_bound = 127 if quant_type == torch.int8 else 7

      m, n, k = 128, 128, 128
      weights = torch.randint(lower_bound, upper_bound, [k, n], dtype=torch.int8, device="cpu")

      packed_weight = self.pack_int4s(weights) if quant_type == torch.quint4x2 else weights
      cuda_weights = self.preprocess_weights_for_mixed_gemm(packed_weight, quant_type).to("cuda")
      weights = weights.to("cuda")

      act = torch.eye(m, dtype=weight_type, device="cuda")
      bias = torch.zeros([n], dtype=weight_type, device='cuda')
      torch_weight_scales = torch.ones_like(bias)
      experts_for_rows = self.setup_experts_for_row(m, 1, 1)

      actual = self.custom_moe_fc(act, experts_for_rows, cuda_weights, torch_weight_scales, bias, "identity")
      torch.allclose(actual, weights, atol=0.0, rtol=0.0)
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
      assert quant_type == torch.int8 or quant_type == torch.quint4x2

      lower_bound = -128 if quant_type == torch.int8 else -8
      upper_bound = 127 if quant_type == torch.int8 else 7

      m, n, k = 128, 128, 128
      weights = torch.randint(lower_bound, upper_bound, [k, n], dtype=torch.int8, device="cpu")

      packed_weight = self.pack_int4s(weights) if quant_type == torch.quint4x2 else weights
      cuda_weights = self.preprocess_weights_for_mixed_gemm(packed_weight, quant_type).to("cuda")
      weights = weights.to("cuda")

      act = torch.eye(m, dtype=weight_type, device="cuda")
      bias = torch.zeros([n], dtype=weight_type, device='cuda')
      torch_weight_scales = torch.ones_like(bias)
      experts_for_rows = self.setup_experts_for_row(m, 1, 1)

      actual = self.custom_moe_fc(act, experts_for_rows, cuda_weights, torch_weight_scales, bias, "identity")
      torch.testing.assert_close(actual, weights, atol=0.0, rtol=0.0, check_dtype=False)
//  [python]
def moe_test_helper(self, dtype, quant_type, rtol, atol, activation_str="gelu", experts_list=[32], hidden_sizes=[1024], inter_sizes=[4096]):
    torch.cuda.empty_cache() # Empty the cache here so a bad ordering does not cause OOM.
    rows = [1, 2, 4]
    ks = range(1, 9)

    for hidden_size in hidden_sizes:
      for inter_size in inter_sizes:
        for experts in experts_list:
          weights = self.generate_weights(hidden_size, inter_size, experts, dtype, quant_type)
          for row in rows:
            for active_rows in [1, row // 2, row]:
              for k in ks:
                if k > experts:
                  continue
                input_dict = self.generate_inputs(row, active_rows, hidden_size, experts, dtype, quant_type)
                input_dict.update(weights)            
                rows_to_check = torch.logical_not(input_dict["finished"])

                # Only take unfinished rows. We can write anything to the output of rows that already complete.
                act_output = self.run_ft_moe(input_dict, row, k, activation_str)[rows_to_check]
                ref_output = self.run_ref_moe(input_dict, k, activation_str)[rows_to_check]

                msg = "Moe Failed on rows={}, active_rows={}, experts={}, k={}, hidden_size={}, inter_size={}" \
                        .format(row, active_rows, experts, k, hidden_size, inter_size)
                t.assert_close(act_output, ref_output, rtol=rtol, atol=atol, msg=msg, check_dtype=False)
//  [python]
def moe_test_helper(self, dtype, quant_type, rtol, atol, activation_str="gelu", experts_list=[32], hidden_sizes=[1024], inter_sizes=[4096]):
    torch.cuda.empty_cache() # Empty the cache here so a bad ordering does not cause OOM.
    rows = [1, 2, 4]
    ks = range(1, 9)

    for hidden_size in hidden_sizes:
      for inter_size in inter_sizes:
        for experts in experts_list:
          weights = self.generate_weights(hidden_size, inter_size, experts, dtype, quant_type)
          for row in rows:
            for active_rows in [1, row // 2, row]:
              for k in ks:
                if k > experts:
                  continue
                input_dict = self.generate_inputs(row, active_rows, hidden_size, experts, dtype, quant_type)
                input_dict.update(weights)            
                rows_to_check = torch.logical_not(input_dict["finished"])

                # Only take unfinished rows. We can write anything to the output of rows that already complete.
                act_output = self.run_ft_moe(input_dict, row, k, activation_str)[rows_to_check]
                ref_output = self.run_ref_moe(input_dict, k, activation_str)[rows_to_check]

                msg = "Moe Failed on rows={}, active_rows={}, experts={}, k={}, hidden_size={}, inter_size={}" \
                        .format(row, active_rows, experts, k, hidden_size, inter_size)
                torch.testing.assert_close(act_output, ref_output, rtol=rtol, atol=atol, msg=msg, check_dtype=False)
//  [python]
def test_pack_unpack_state(self):
        sm = TestScript.DerivedStateModule()
        x = torch.rand(3, 4, dtype=torch.float)
        self.assertEqual(sm(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))

        # Test save path
        self.assertFalse(sm.pack_called.item())
        self.assertFalse(sm.unpack_called.item())
        imported = self.getExportImportCopyWithPacking(sm)
        # ensure pack was called before serialization
        self.assertTrue(sm.pack_called.item())
        # ensure unpack was called after serialization so as to leave the module in an initialized state
        self.assertTrue(sm.unpack_called.item())

        self.assertEqual(sm.derived, torch.neg(sm.param))

        # Test load paths
        self.assertTrue(imported.unpack_called.item())
        self.assertEqual(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
//  [python]
def test_torchscript_multi_head_attn_fast_path(self):
        src_l = 3
        bsz = 5
        embed_size = 8
        nhead = 2
        multi_head_attn = torch.nn.MultiheadAttention(embed_size, nhead, batch_first=True)
        multi_head_attn = multi_head_attn.eval()

        query = key = value = torch.rand((bsz, src_l, embed_size))

        with torch.no_grad():
            py_out = multi_head_attn(query, key, value)
            mha = torch.jit.script(multi_head_attn)
            jit_out = mha(query, key, value)
        torch.testing.assert_allclose(jit_out, py_out)
//  [python]
def test_torchscript_multi_head_attn_fast_path(self):
        src_l = 3
        bsz = 5
        embed_size = 8
        nhead = 2
        multi_head_attn = torch.nn.MultiheadAttention(embed_size, nhead, batch_first=True)
        multi_head_attn = multi_head_attn.eval()

        query = key = value = torch.rand((bsz, src_l, embed_size))

        with torch.no_grad():
            py_out = multi_head_attn(query, key, value)
            mha = torch.jit.script(multi_head_attn)
            jit_out = mha(query, key, value)
        torch.testing.assert_close(jit_out, py_out)
//  [python]
def test_apply_rotary(self, q_states, k_states, cached_cos, cached_sin,
                          position_ids_1d, gt):
        q_embed, k_embed = apply_rotary_pos_emb(q_states, k_states, cached_cos,
                                                cached_sin, None,
                                                position_ids_1d)
        q_gt, k_gt = gt

        rtol = None
        atol = None
        if q_states.dtype == torch.float16:
            rtol = 1e-5
            atol = 1e-3
        torch.testing.assert_allclose(q_embed, q_gt, rtol=rtol, atol=atol)
        torch.testing.assert_allclose(k_embed, k_gt, rtol=rtol, atol=atol)
//  [python]
def test_apply_rotary(self, q_states, k_states, cached_cos, cached_sin,
                          position_ids_1d, gt):
        q_embed, k_embed = apply_rotary_pos_emb(q_states, k_states, cached_cos,
                                                cached_sin, None,
                                                position_ids_1d)
        q_gt, k_gt = gt

        rtol = None
        atol = None
        if q_states.dtype == torch.float16:
            rtol = 1e-5
            atol = 1e-3
        torch.testing.assert_close(q_embed, q_gt, rtol=rtol, atol=atol)
        torch.testing.assert_close(k_embed, k_gt, rtol=rtol, atol=atol)
//  [python]
def test_dataset_categorical():
    """Test the processing of categorical data."""
    key = "foo"
    key2 = "bar"
    batch_size = 321
    data = TensorDict(
        {
            key: np.random.binomial(1, 0.5, batch_size)[:, None].astype(np.int32),
            key2: np.random.rand(batch_size)[:, None].astype(np.float32),
        },
        batch_size=batch_size,
    )

    assert data[key].shape == (batch_size, 1)
    assert data[key2].shape == (batch_size, 1)

    # should be one hot encoded
    dataset = convert_one_hot(data, one_hot_sizes={key: 2})
    assert dataset[key].shape == (batch_size, 2)
    assert dataset[key2].shape == (batch_size, 1)

    dataset = convert_one_hot(data, one_hot_sizes={key: 3})
    assert dataset[key].shape == (batch_size, 3)
    assert dataset[key2].shape == (batch_size, 1)
    torch.testing.assert_allclose(dataset[key][:, -1], torch.zeros_like(dataset[key][:, -1]))
//  [python]
def test_dataset_categorical():
    """Test the processing of categorical data."""
    key = "foo"
    key2 = "bar"
    batch_size = 321
    data = TensorDict(
        {
            key: np.random.binomial(1, 0.5, batch_size)[:, None].astype(np.int32),
            key2: np.random.rand(batch_size)[:, None].astype(np.float32),
        },
        batch_size=batch_size,
    )

    assert data[key].shape == (batch_size, 1)
    assert data[key2].shape == (batch_size, 1)

    # should be one hot encoded
    dataset = convert_one_hot(data, one_hot_sizes={key: 2})
    assert dataset[key].shape == (batch_size, 2)
    assert dataset[key2].shape == (batch_size, 1)

    dataset = convert_one_hot(data, one_hot_sizes={key: 3})
    assert dataset[key].shape == (batch_size, 3)
    assert dataset[key2].shape == (batch_size, 1)
    torch.testing.assert_close(dataset[key][:, -1], torch.zeros_like(dataset[key][:, -1]))
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

        class M(torch.nn.Module):
            def __init__(self):
                super(M, self).__init__()
                self.weight_0 = torch.nn.Parameter(torch.rand(weight_0_shape))
                self.weight_1 = torch.nn.Parameter(torch.rand(weight_1_shape))

            def forward(self, x):
                o = F.linear(x, self.weight_0)
                o = F.dropout(o, training=self.training)
                o = F.linear(o, self.weight_1)
                return o

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        self.assertEqual(ref_res, res)
//  [python]
def test_config_flag_is_respected(self):
        def split_with_cat(x):
            fs = torch.split(x, [4, 4, 24], dim=-1)
            item0 = fs[0]
            item1 = fs[1]
            item2 = fs[2]

            final_items = [item0, item1]
            final_items.extend(item2.split((4, 4, 4, 4, 4, 4), 1))

            return torch.cat(final_items, dim=1)

        args = [
            torch.randn(2, 32),
        ]

        expected = split_with_cat(*args)
        actual = torch.compile(split_with_cat)(*args)

        self.assertEqual(torch.allclose(actual, expected),True)
        self.assertEqual(
            counters["inductor"]["merge_splits_pass"],
            0,
        )
        self.assertEqual(
            counters["inductor"]["normalization_pass"],
            0,
        )