jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_dynamic_skip_quantization(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_skip_quantization(layer_names={"conv2"})
        q_util.quantize_module(module)
        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))
//  [python]
def test_gather(self, input, start_loc, seq_lens, adapter_ids, ranks,
                    world_size, gt, valid_mask):
        max_seq_len = max(seq_lens)
        output = rearange_all_gather(input,
                                     start_loc,
                                     seq_lens,
                                     adapter_ids,
                                     ranks,
                                     world_size,
                                     max_seq_len=max_seq_len)
        output = output.where(valid_mask, output.new_tensor(0))
        gt = gt.where(valid_mask, gt.new_tensor(0))
        torch.testing.assert_allclose(output, gt)
//  [python]
def test_gather(self, input, start_loc, seq_lens, adapter_ids, ranks,
                    world_size, gt, valid_mask):
        max_seq_len = max(seq_lens)
        output = rearange_all_gather(input,
                                     start_loc,
                                     seq_lens,
                                     adapter_ids,
                                     ranks,
                                     world_size,
                                     max_seq_len=max_seq_len)
        output = output.where(valid_mask, output.new_tensor(0))
        gt = gt.where(valid_mask, gt.new_tensor(0))
        torch.testing.assert_close(output, gt)
//  [python]
def test_nested_module_replacement(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

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

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))
//  [python]
def test_nested_module_replacement(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

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

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))
//  [python]
def test_non_default_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

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

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer(default_quant_modules_calibrator_inputs="max", default_quant_modules_calibrator_weights="max")
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_non_default_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

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

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer(default_quant_modules_calibrator_inputs="max", default_quant_modules_calibrator_weights="max")
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_custom_quantized_mapping_register_with_decorator(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

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

        @register_quantized_module(float_source=MyBlock)
        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

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

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, n_clss))
        self.assertTrue(MyQuantizedBlock is not None)
        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))
//  [python]
def test_custom_quantized_mapping_register_with_decorator(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

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

        @register_quantized_module(float_source=MyBlock)
        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

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

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))
        self.assertTrue(MyQuantizedBlock is not None)
        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))
//  [python]
def test_global_blocked_params(self) -> None:
        expected_global_params = (
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, dtype=torch.float),
        )
        torch.testing.assert_allclose(
            self._distributor.global_blocked_params,
            expected_global_params,
        )
//  [python]
def test_global_blocked_params(self) -> None:
        expected_global_params = (
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, dtype=torch.float),
        )
        torch.testing.assert_close(
            self._distributor.global_blocked_params,
            expected_global_params,
        )
//  [python]
def test_different_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_quantization_mapping(
            layer_names={"conv1"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="max"),
            weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
        )
        q_util.register_quantization_mapping(
            layer_names={"conv2"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
            weights_quant_descriptor=QuantDescriptor(calib_method="max"),
        )
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            self.assertEqual(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._input_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_different_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_quantization_mapping(
            layer_names={"conv1"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="max"),
            weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
        )
        q_util.register_quantization_mapping(
            layer_names={"conv2"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
            weights_quant_descriptor=QuantDescriptor(calib_method="max"),
        )
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._input_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._weight_quantizer._calibrator) == MaxCalibrator)
//  [python]
def test_trace_converage(self, id_or_path: str) -> None:
        model, info = create_model(id_or_path=id_or_path, torchscript=True)
        traced = trace_model(model, info)
        golden_outputs = model(**info.example_inputs)
        inputs = _kwargs_to_args(info.input_order, **info.example_inputs)
        traced_outputs = traced(*inputs)
        self.assertEqual(len(golden_outputs), len(traced_outputs))
        for a, b in zip(golden_outputs, traced_outputs):
            self.assertEqual(a.shape,b.shape)
            t.assert_allclose(a, b, rtol=0.01, atol=0.01)
//  [python]
def test_trace_converage(self, id_or_path: str) -> None:
        model, info = create_model(id_or_path=id_or_path, torchscript=True)
        traced = trace_model(model, info)
        golden_outputs = model(**info.example_inputs)
        inputs = _kwargs_to_args(info.input_order, **info.example_inputs)
        traced_outputs = traced(*inputs)
        self.assertEqual(len(golden_outputs), len(traced_outputs))
        for a, b in zip(golden_outputs, traced_outputs):
            torch.testing.assert_close(a, b, rtol=0.01, atol=0.01)
//  [python]
def _test_code_common(
        self,
        mod,
        inputs,
        include_ops,
        exclude_ops,
        atol=1e-5,
        rtol=1.3e-6,
        check_quantization=False,
        check_dynamic=None,
    ):
        with torch.no_grad():
            clone_inputs = self._clone_inputs(inputs)
            if check_quantization:
                mod = self._generate_qdq_quantized_model(mod, inputs)
            expected = mod(*inputs)
            actual, (source_code,) = run_and_get_code(
                torch.compile(mod, fullgraph=True, dynamic=check_dynamic),
                *clone_inputs,
            )
            for op in include_ops:
                self.assertIn(op, source_code)
            for op in exclude_ops:
                self.assertNotIn(op, source_code)
            if check_dynamic is not None:
                _check_has_dynamic_shape(self, source_code)
            if not check_quantization:
                # Skip due to reduce range setting for Quantization on preCI system.
                torch.testing.assert_allclose(actual, expected, atol=atol, rtol=rtol)
//  [python]
def _test_code_common(
        self,
        mod,
        inputs,
        include_ops,
        exclude_ops,
        atol=1e-5,
        rtol=1.3e-6,
        check_quantization=False,
        check_dynamic=None,
    ):
        with torch.no_grad():
            clone_inputs = self._clone_inputs(inputs)
            if check_quantization:
                mod = self._generate_qdq_quantized_model(mod, inputs)
            expected = mod(*inputs)
            actual, (source_code,) = run_and_get_code(
                torch.compile(mod, fullgraph=True, dynamic=check_dynamic),
                *clone_inputs,
            )
            for op in include_ops:
                self.assertIn(op, source_code)
            for op in exclude_ops:
                self.assertNotIn(op, source_code)
            if check_dynamic is not None:
                _check_has_dynamic_shape(self, source_code)
            if not check_quantization:
                # Skip due to reduce range setting for Quantization on preCI system.
                torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
//  [python]
def test_save_load_params(self, gp_fit, tmpdir):
        gp2 = clone(gp_fit).initialize()

        # check first that parameters are not equal
        for (_, p0), (_, p1) in zip(
                gp_fit.get_all_learnable_params(), gp2.get_all_learnable_params(),
        ):
            assert not (p0 == p1).all()

        # save and load params to gp2
        p_module = tmpdir.join('module.pt')
        p_likelihood = tmpdir.join('likelihood.pt')
        with open(str(p_module), 'wb') as fm, open(str(p_likelihood), 'wb') as fll:
            gp_fit.save_params(f_params=fm, f_likelihood=fll)
        with open(str(p_module), 'rb') as fm, open(str(p_likelihood), 'rb') as fll:
            gp2.load_params(f_params=fm, f_likelihood=fll)

        # now parameters should be equal
        for (n0, p0), (n1, p1) in zip(
                gp_fit.get_all_learnable_params(), gp2.get_all_learnable_params(),
        ):
            assert n0 == n1
            t.assert_allclose(p0, p1)
//  [python]
def test_save_load_params(self, gp_fit, tmpdir):
        gp2 = clone(gp_fit).initialize()

        # check first that parameters are not equal
        for (_, p0), (_, p1) in zip(
                gp_fit.get_all_learnable_params(), gp2.get_all_learnable_params(),
        ):
            assert not (p0 == p1).all()

        # save and load params to gp2
        p_module = tmpdir.join('module.pt')
        p_likelihood = tmpdir.join('likelihood.pt')
        with open(str(p_module), 'wb') as fm, open(str(p_likelihood), 'wb') as fll:
            gp_fit.save_params(f_params=fm, f_likelihood=fll)
        with open(str(p_module), 'rb') as fm, open(str(p_likelihood), 'rb') as fll:
            gp2.load_params(f_params=fm, f_likelihood=fll)

        # now parameters should be equal
        for (n0, p0), (n1, p1) in zip(
                gp_fit.get_all_learnable_params(), gp2.get_all_learnable_params(),
        ):
            assert n0 == n1
            torch.testing.assert_close(p0, p1)
//  [python]
def test_update_params(self) -> None:
        # Explicitly disable the gradient of the bias layer and call merge_and_block_gradients()
        # to update the local gradient selector.
        self._model[0].weight.grad = torch.ones((5, 10))
        self._model[0].bias.grad = None
        self._distributor.merge_and_block_gradients()

        actual_masked_blocked_params = self._distributor.local_masked_blocked_params

        masked_blocked_search_directions = (
            torch.arange(5 * 5, dtype=torch.float).reshape(5, 5),
            torch.arange(5 * 5, 2 * 5 * 5, dtype=torch.float).reshape(5, 5),
        )
        self._distributor.update_params(
            masked_blocked_search_directions,
        )

        expected_masked_blocked_params = (
            torch.arange(5 * 5, dtype=torch.float).reshape(5, 5),
            torch.arange(5 * 5, 2 * 5 * 5, dtype=torch.float).reshape(5, 5),
        )
        torch.testing.assert_allclose(
            actual_masked_blocked_params, expected_masked_blocked_params
        )
//  [python]
def test_update_params(self) -> None:
        # Explicitly disable the gradient of the bias layer and call merge_and_block_gradients()
        # to update the local gradient selector.
        self._model[0].weight.grad = torch.ones((5, 10))
        self._model[0].bias.grad = None
        self._distributor.merge_and_block_gradients()

        actual_masked_blocked_params = self._distributor.local_masked_blocked_params

        masked_blocked_search_directions = (
            torch.arange(5 * 5, dtype=torch.float).reshape(5, 5),
            torch.arange(5 * 5, 2 * 5 * 5, dtype=torch.float).reshape(5, 5),
        )
        self._distributor.update_params(
            masked_blocked_search_directions,
        )

        expected_masked_blocked_params = (
            torch.arange(5 * 5, dtype=torch.float).reshape(5, 5),
            torch.arange(5 * 5, 2 * 5 * 5, dtype=torch.float).reshape(5, 5),
        )
        torch.testing.assert_close(
            actual_masked_blocked_params, expected_masked_blocked_params
        )
//  [python]
def test_local_blocked_params(self) -> None:
        # In Distributor, because there is no global vs. local boundary concept,
        # globl and local blocked params are always identical.
        expected_local_params = (
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, dtype=torch.float),
        )
        self.assertEqual(self._distributor.local_blocked_params, expected_local_params)
//  [python]
def test_local_blocked_params(self) -> None:
        # In Distributor, because there is no global vs. local boundary concept,
        # globl and local blocked params are always identical.
        expected_local_params = (
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, 5, dtype=torch.float),
            torch.zeros(5, dtype=torch.float),
        )
        torch.testing.assert_close(
            self._distributor.local_blocked_params,
            expected_local_params,
        )
//  [python]
def test_non_removable_id(
        self,
        test_name="non_removable_id",
        dtype="float16",
    ):
        # make a graph like below:
        # x0 = tensor
        # y = id(x0)
        batch_sizes = [1, self.BATCH_SIZE]
        batch_dim = shape_utils.gen_int_var_min_max(batch_sizes, "batch_0")
        M = 10
        X0 = gen_input_tensor([batch_dim, IntImm(M)], name="x0", dtype=dtype)
        Y = ops.identity()(X0)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        target = detect_target()
        dll_name = f"test_{self._test_id}.so"
        module = compile_model(Y, target, "./tmp", test_name, dll_name=dll_name)
        self._test_id += 1

        sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
        self.assertEqual(len(sorted_ops), 1)
        self.assertEqual(sorted_ops[0]._attrs["op"], "identity")

        for batch in [1, self.BATCH_SIZE]:
            x0_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = x0_pt

            y = get_torch_empty_tensor(y_pt.size(), dtype)
            inputs = {"x0": x0_pt}
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(y_pt, y, atol=0.01, rtol=0.01)
//  [python]
def test_non_removable_id(
        self,
        test_name="non_removable_id",
        dtype="float16",
    ):
        # make a graph like below:
        # x0 = tensor
        # y = id(x0)
        batch_sizes = [1, self.BATCH_SIZE]
        batch_dim = shape_utils.gen_int_var_min_max(batch_sizes, "batch_0")
        M = 10
        X0 = gen_input_tensor([batch_dim, IntImm(M)], name="x0", dtype=dtype)
        Y = ops.identity()(X0)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True

        target = detect_target()
        dll_name = f"test_{self._test_id}.so"
        module = compile_model(Y, target, "./tmp", test_name, dll_name=dll_name)
        self._test_id += 1

        sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
        self.assertEqual(len(sorted_ops), 1)
        self.assertEqual(sorted_ops[0]._attrs["op"], "identity")

        for batch in [1, self.BATCH_SIZE]:
            x0_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = x0_pt

            y = get_torch_empty_tensor(y_pt.size(), dtype)
            inputs = {"x0": x0_pt}
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y, atol=0.01, rtol=0.01)
//  [python]
def test_merge_and_block_gradients(self) -> None:
        self._model[0].weight.grad = torch.ones((5, 10))
        self._model[0].bias.grad = None
        actual_local_masked_block_grads = self._distributor.merge_and_block_gradients()
        expected_local_masked_block_grads = (
            torch.ones((5, 5)),
            torch.ones((5, 5)),
        )
        torch.testing.assert_allclose(
            actual_local_masked_block_grads, expected_local_masked_block_grads
        )
//  [python]
def test_merge_and_block_gradients(self) -> None:
        self._model[0].weight.grad = torch.ones((5, 10))
        self._model[0].bias.grad = None
        actual_local_masked_block_grads = self._distributor.merge_and_block_gradients()
        expected_local_masked_block_grads = (
            torch.ones((5, 5)),
            torch.ones((5, 5)),
        )
        torch.testing.assert_close(
            actual_local_masked_block_grads, expected_local_masked_block_grads
        )
//  [python]
def test_concrete_arg_none_assert(self):
        class Foo(torch.nn.Module):
            def forward(self, x, val=None):
                return x if val is None else x + val

        f = Foo()
        traced = torch.fx.symbolic_trace(f, concrete_args={'val' : None})
        with self.assertRaisesRegex(AssertionError, 'val has been specialized to have value None'):
            traced(torch.randn(5), torch.randn(5))

        x = torch.randn(5)
        self.assertEqual(traced(x), f(x))
//  [python]
def test_concrete_arg_none_assert(self):
        class Foo(torch.nn.Module):
            def forward(self, x, val=None):
                return x if val is None else x + val

        f = Foo()
        traced = torch.fx.symbolic_trace(f, concrete_args={'val' : None})
        with self.assertRaisesRegex(AssertionError, 'val has been specialized to have value None'):
            traced(torch.randn(5), torch.randn(5))

        x = torch.randn(5)
        torch.testing.assert_close(traced(x), f(x))
//  [python]
def test_tuple_no_subscript(self):
        def foo(x : Tuple):
            return x[0]

        traced = torch.fx.symbolic_trace(foo)
        x = (torch.randn(5, 3),)
        t.assert_close(traced(x), x[0])

        bio = io.BytesIO()

        torch.save(traced, bio)

        bio.seek(0)

        loaded = torch.load(bio)

        t.assert_close(loaded(x), x[0])
//  [python]
def test_tuple_no_subscript(self):
        def foo(x : Tuple):
            return x[0]

        traced = torch.fx.symbolic_trace(foo)
        x = (torch.randn(5, 3),)
        torch.testing.assert_close(traced(x), x[0])

        bio = io.BytesIO()

        torch.save(traced, bio)

        bio.seek(0)

        loaded = torch.load(bio)

        torch.testing.assert_close(loaded(x), x[0])
//  [python]
def test_interpreter_default_args(self):
        class Model(torch.nn.Module):
            def forward(self, x, y=3.14159):
                return x + y

        model = Model()
        gm = torch.fx.symbolic_trace(model)

        interp = Interpreter(gm)
        x = torch.randn(5, 3)
        out = interp.run(x)
        torch.testing.assert_allclose(out, x + 3.14159)
//  [python]
def test_interpreter_default_args(self):
        class Model(torch.nn.Module):
            def forward(self, x, y=3.14159):
                return x + y

        model = Model()
        gm = torch.fx.symbolic_trace(model)

        interp = Interpreter(gm)
        x = torch.randn(5, 3)
        out = interp.run(x)
        torch.testing.assert_close(out, x + 3.14159)
//  [python]
def test_trace_multiple_funcs(self):
        class Foo(torch.nn.Module):
            def forward(self, x, y):
                return x + y

            def minus_forward(self, x, y):
                return x - y

            def multiply_forward(self, x, y):
                return x * y

        f = Foo()
        x, y = torch.randn(5), torch.randn(5)

        print(torch.__version__)

        tracer = Tracer()
        t.assert_close(GraphModule(f, tracer.trace(f))(x, y), f(x, y))

        tracer.traced_func_name = "minus_forward"
        t.assert_close(
            GraphModule(f, tracer.trace(f))(x, y),
            f.minus_forward(x, y),
        )

        tracer.traced_func_name = "multiply_forward"
        t.assert_close(
            GraphModule(f, tracer.trace(f))(x, y),
            f.multiply_forward(x, y),
        )

        tracer.traced_func_name = "add_forward"
        with self.assertRaisesRegex(AssertionError, "doesn't exist in"):
            tracer.trace(f)
//  [python]
def test_trace_multiple_funcs(self):
        class Foo(torch.nn.Module):
            def forward(self, x, y):
                return x + y

            def minus_forward(self, x, y):
                return x - y

            def multiply_forward(self, x, y):
                return x * y

        f = Foo()
        x, y = torch.randn(5), torch.randn(5)

        print(torch.__version__)

        tracer = Tracer()
        torch.testing.assert_close(GraphModule(f, tracer.trace(f))(x, y), f(x, y))

        tracer.traced_func_name = "minus_forward"
        torch.testing.assert_close(
            GraphModule(f, tracer.trace(f))(x, y),
            f.minus_forward(x, y),
        )

        tracer.traced_func_name = "multiply_forward"
        torch.testing.assert_close(
            GraphModule(f, tracer.trace(f))(x, y),
            f.multiply_forward(x, y),
        )

        tracer.traced_func_name = "add_forward"
        with self.assertRaisesRegex(AssertionError, "doesn't exist in"):
            tracer.trace(f)
//  [python]
def test_gemm_rrr_to_rcr(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        x = Tensor([32, 52], dtype=dtype, is_input=True, name="x")
        w = Tensor([32, 52], dtype=dtype, is_input=True, name="w")
        y = ops.permute()(x, dims=[1, 0])
        z = ops.gemm_rrr()(w, y)
        z._attrs["is_output"] = True
        z._attrs["name"] = "z"

        module = compile_model(z, target, "./tmp", f"test_gemm_rrr_to_rcr_{dtype}")
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "permute"))
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rrr"))
        self.assertTrue(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rcr"))

        x_pt = get_random_torch_tensor([32, 52], dtype)
        w_pt = get_random_torch_tensor([32, 52], dtype)
        y_pt = x_pt.t()
        z_pt = torch.matmul(w_pt, y_pt)
        z_ait = torch.empty_like(z_pt)
        module.run_with_tensors({"x": x_pt, "w": w_pt}, {"z": z_ait})

        torch.testing.assert_allclose(z_ait, z_pt, atol=1e-1, rtol=1e-1)
//  [python]
def test_gemm_rrr_to_rcr(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        x = Tensor([32, 52], dtype=dtype, is_input=True, name="x")
        w = Tensor([32, 52], dtype=dtype, is_input=True, name="w")
        y = ops.permute()(x, dims=[1, 0])
        z = ops.gemm_rrr()(w, y)
        z._attrs["is_output"] = True
        z._attrs["name"] = "z"

        module = compile_model(z, target, "./tmp", f"test_gemm_rrr_to_rcr_{dtype}")
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "permute"))
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rrr"))
        self.assertTrue(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rcr"))

        x_pt = get_random_torch_tensor([32, 52], dtype)
        w_pt = get_random_torch_tensor([32, 52], dtype)
        y_pt = x_pt.t()
        z_pt = torch.matmul(w_pt, y_pt)
        z_ait = torch.empty_like(z_pt)
        module.run_with_tensors({"x": x_pt, "w": w_pt}, {"z": z_ait})

        torch.testing.assert_close(z_ait, z_pt, atol=1e-1, rtol=1e-1)
//  [python]
def test_trace_buffer_slice(self):
        bs, d_hid = 10, 23

        class ExampleCode(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mm_param = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.mm_param2 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.lin = torch.nn.Linear(d_hid, d_hid)
                self.register_buffer('buffer', torch.randn(bs + 100, d_hid))

            def forward(self, x):
                x = torch.mm(x, self.mm_param)
                skip_connection = x
                x = torch.relu(x)
                x = torch.mm(x, self.mm_param) + self.buffer[:x.shape[0]]
                x = self.lin(x)
                x = torch.relu(x)
                x = x + skip_connection
                x = torch.mm(x, self.mm_param2)
                x = self.lin(x)
                return x


        ec = ExampleCode()

        traced = torch.fx.symbolic_trace(ec)

        x = torch.randn(bs, d_hid)
        torch.testing.assert_allclose(ec(x), traced(x))
//  [python]
def test_trace_buffer_slice(self):
        bs, d_hid = 10, 23

        class ExampleCode(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mm_param = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.mm_param2 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
                self.lin = torch.nn.Linear(d_hid, d_hid)
                self.register_buffer('buffer', torch.randn(bs + 100, d_hid))

            def forward(self, x):
                x = torch.mm(x, self.mm_param)
                skip_connection = x
                x = torch.relu(x)
                x = torch.mm(x, self.mm_param) + self.buffer[:x.shape[0]]
                x = self.lin(x)
                x = torch.relu(x)
                x = x + skip_connection
                x = torch.mm(x, self.mm_param2)
                x = self.lin(x)
                return x


        ec = ExampleCode()

        traced = torch.fx.symbolic_trace(ec)

        x = torch.randn(bs, d_hid)
        torch.testing.assert_close(ec(x), traced(x))
//  [python]
def test_gemm_rcr_to_rrr(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        x = Tensor([52, 32], dtype=dtype, is_input=True, name="x")
        w = Tensor([32, 52], dtype=dtype, is_input=True, name="w")
        y = ops.permute()(x, dims=[1, 0])
        z = ops.gemm_rcr()(w, y)
        z._attrs["is_output"] = True
        z._attrs["name"] = "z"

        module = compile_model(
            z,
            target,
            "./tmp",
            f"test_gemm_rcr_to_rrr_{dtype}",
        )
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "permute"))
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rcr"))
        self.assertTrue(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rrr"))

        x_pt = get_random_torch_tensor([52, 32], dtype)
        w_pt = get_random_torch_tensor([32, 52], dtype)
        z_pt = torch.matmul(w_pt, x_pt)
        z_ait = torch.empty_like(z_pt)
        module.run_with_tensors({"x": x_pt, "w": w_pt}, {"z": z_ait})

        torch.testing.assert_allclose(z_ait, z_pt, atol=1e-1, rtol=1e-1)
//  [python]
def test_gemm_rcr_to_rrr(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        x = Tensor([52, 32], dtype=dtype, is_input=True, name="x")
        w = Tensor([32, 52], dtype=dtype, is_input=True, name="w")
        y = ops.permute()(x, dims=[1, 0])
        z = ops.gemm_rcr()(w, y)
        z._attrs["is_output"] = True
        z._attrs["name"] = "z"

        module = compile_model(
            z,
            target,
            "./tmp",
            f"test_gemm_rcr_to_rrr_{dtype}",
        )
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "permute"))
        self.assertFalse(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rcr"))
        self.assertTrue(test_utils.graph_has_op(module.debug_sorted_graph, "gemm_rrr"))

        x_pt = get_random_torch_tensor([52, 32], dtype)
        w_pt = get_random_torch_tensor([32, 52], dtype)
        z_pt = torch.matmul(w_pt, x_pt)
        z_ait = torch.empty_like(z_pt)
        module.run_with_tensors({"x": x_pt, "w": w_pt}, {"z": z_ait})

        torch.testing.assert_close(z_ait, z_pt, atol=1e-1, rtol=1e-1)
//  [python]
def test_module_deepcopy_edit_nodes(self):
        class Foo(torch.nn.Module):
            def forward(self, x):
                return torch.relu(x)

        traced1 = symbolic_trace(Foo())
        copied = copy.deepcopy(traced1)

        for node in copied.graph.nodes:
            if node.target == torch.relu:
                node.target = torch.neg

        copied.recompile()
        traced1.recompile()

        x = torch.randn(15, 15)
        self.assertEqual(traced1(x), torch.relu(x))
        self.assertEqual(copied(x), torch.neg(x))
//  [python]
def test_module_deepcopy_edit_nodes(self):
        class Foo(torch.nn.Module):
            def forward(self, x):
                return torch.relu(x)

        traced1 = symbolic_trace(Foo())
        copied = copy.deepcopy(traced1)

        for node in copied.graph.nodes:
            if node.target == torch.relu:
                node.target = torch.neg

        copied.recompile()
        traced1.recompile()

        x = torch.randn(15, 15)
        torch.testing.assert_close(traced1(x), torch.relu(x))
        torch.testing.assert_close(copied(x), torch.neg(x))
//  [python]
def test_update_weights(use_async):
    def create_env():
        return ContinuousActionVecMockEnv()

    n_actions = ContinuousActionVecMockEnv().action_spec.shape[-1]
    policy = SafeModule(
        torch.nn.LazyLinear(n_actions), in_keys=["observation"], out_keys=["action"]
    )
    policy(create_env().reset())

    collector_class = (
        MultiSyncDataCollector if not use_async else MultiaSyncDataCollector
    )
    collector = collector_class(
        [create_env] * 3,
        policy=policy,
        device=[torch.device("cuda:0")] * 3,
        storing_device=[torch.device("cuda:0")] * 3,
        frames_per_batch=20,
        cat_results="stack",
    )
    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            t.assert_close(
                state_dict[f"worker{worker}"]["policy_state_dict"][k],
                policy_state_dict[k].cpu(),
            )

    # change policy weights
    for p in policy.parameters():
        p.data += torch.randn_like(p)

    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    # check they don't match
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            with pytest.raises(AssertionError):
                t.assert_close(
                    state_dict[f"worker{worker}"]["policy_state_dict"][k],
                    policy_state_dict[k].cpu(),
                )

    # update weights
    collector.update_policy_weights_()

    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            t.assert_close(
                state_dict[f"worker{worker}"]["policy_state_dict"][k],
                policy_state_dict[k].cpu(),
            )

    collector.shutdown()
    del collector
//  [python]
def test_update_weights(use_async):
    def create_env():
        return ContinuousActionVecMockEnv()

    n_actions = ContinuousActionVecMockEnv().action_spec.shape[-1]
    policy = SafeModule(
        torch.nn.LazyLinear(n_actions), in_keys=["observation"], out_keys=["action"]
    )
    policy(create_env().reset())

    collector_class = (
        MultiSyncDataCollector if not use_async else MultiaSyncDataCollector
    )
    collector = collector_class(
        [create_env] * 3,
        policy=policy,
        device=[torch.device("cuda:0")] * 3,
        storing_device=[torch.device("cuda:0")] * 3,
        frames_per_batch=20,
        cat_results="stack",
    )
    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            torch.testing.assert_close(
                state_dict[f"worker{worker}"]["policy_state_dict"][k],
                policy_state_dict[k].cpu(),
            )

    # change policy weights
    for p in policy.parameters():
        p.data += torch.randn_like(p)

    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    # check they don't match
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            with pytest.raises(AssertionError):
                torch.testing.assert_close(
                    state_dict[f"worker{worker}"]["policy_state_dict"][k],
                    policy_state_dict[k].cpu(),
                )

    # update weights
    collector.update_policy_weights_()

    # collect state_dict
    state_dict = collector.state_dict()
    policy_state_dict = policy.state_dict()
    for worker in range(3):
        for k in state_dict[f"worker{worker}"]["policy_state_dict"]:
            torch.testing.assert_close(
                state_dict[f"worker{worker}"]["policy_state_dict"][k],
                policy_state_dict[k].cpu(),
            )

    collector.shutdown()
    del collector
//  [python]
def test_graph_module_replicate_for_dp(self):
        class Foo(torch.nn.Module):
            def forward(self, x):
                return torch.relu(x)

        gm = torch.fx.symbolic_trace(Foo())

        x = torch.randn(5, 3)
        out = gm(x)

        replica = gm._replicate_for_data_parallel()
        out_replica = replica(x)

        self.assertEqual(out_replica, out)
//  [python]
def test_graph_module_replicate_for_dp(self):
        class Foo(torch.nn.Module):
            def forward(self, x):
                return torch.relu(x)

        gm = torch.fx.symbolic_trace(Foo())

        x = torch.randn(5, 3)
        out = gm(x)

        replica = gm._replicate_for_data_parallel()
        out_replica = replica(x)

        torch.testing.assert_close(out_replica, out)
//  [python]
def check_outputs(new_output, original_output):
        t.assert_allclose(new_output, torch.tensor([5]))
//  [python]
def check_outputs(new_output, original_output):
        torch.testing.assert_close(new_output, torch.tensor([5]))
//  [python]
def test_geglu(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = GEGLU(in_features, out_features,
                  bias=bias).cuda().to(dtype=dtype).eval()
        x = torch.randn(N, in_features).cuda().to(dtype=dtype)
        out = m(x)
        out_opt = m(x, enable_opt=True)

        torch.testing.assert_allclose(out_opt, out, rtol=2e-2, atol=2e-2)
//  [python]
def test_geglu(dtype, bias, in_features, out_features, N):
    with torch.no_grad():
        m = GEGLU(in_features, out_features,
                  bias=bias).cuda().to(dtype=dtype).eval()
        x = torch.randn(N, in_features).cuda().to(dtype=dtype)
        out = m(x)
        out_opt = m(x, enable_opt=True)

        torch.testing.assert_close(out_opt, out, rtol=2e-2, atol=2e-2)
//  [python]
def test_load_mobile_module(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, x):
                return x + 10

        input = torch.tensor([1])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        self.assertEqual(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        self.assertEqual(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        self.assertEqual(script_module_result, mobile_module_run_method_result)
//  [python]
def test_load_mobile_module(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, x):
                return x + 10

        input = torch.tensor([1])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        torch.testing.assert_close(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_close(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_close(
            script_module_result, mobile_module_run_method_result
        )
//  [python]
def test_delete_unused_submodules_leaf(self):
        class SubModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(10, 10)
                self.relu = torch.nn.ReLU()

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

        class Model(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.submod = SubModule()

            def forward(self, x):
                x = self.submod(x)
                return x

        model = Model()

        class MyCustomTracer(torch.fx.Tracer):
            def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> bool:
                return module_qualified_name == "submod"

        inputs = torch.randn(1, 10)
        traced_graph = MyCustomTracer().trace(model)
        gm2 = torch.fx.GraphModule(model, traced_graph)
        gm2.delete_all_unused_submodules()
        torch.testing.assert_allclose(gm2(inputs), model(inputs))