jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_supported_input_structure(self, T):

        batch_size = 32
        num_classes = 100

        dataset = self.DummyDataset(size=batch_size, num_classes=num_classes)

        cutmix_mixup = T(num_classes=num_classes)

        dl = DataLoader(dataset, batch_size=batch_size)

        # Input sanity checks
        img, target = next(iter(dl))
        input_img_size = img.shape[-3:]
        assert isinstance(img, torch.Tensor) and isinstance(target, torch.Tensor)
        assert target.shape == (batch_size,)

        def check_output(img, target):
            assert img.shape == (batch_size, *input_img_size)
            assert target.shape == (batch_size, num_classes)
            torch.testing.assert_allclose(target.sum(axis=-1), torch.ones(batch_size))
            num_non_zero_labels = (target != 0).sum(axis=-1)
            assert (num_non_zero_labels == 2).all()

        # After Dataloader, as unpacked input
        img, target = next(iter(dl))
        assert target.shape == (batch_size,)
        img, target = cutmix_mixup(img, target)
        check_output(img, target)

        # After Dataloader, as packed input
        packed_from_dl = next(iter(dl))
        assert isinstance(packed_from_dl, list)
        img, target = cutmix_mixup(packed_from_dl)
        check_output(img, target)

        # As collation function. We expect default_collate to be used by users.
        def collate_fn_1(batch):
            return cutmix_mixup(default_collate(batch))

        def collate_fn_2(batch):
            return cutmix_mixup(*default_collate(batch))

        for collate_fn in (collate_fn_1, collate_fn_2):
            dl = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)
            img, target = next(iter(dl))
            check_output(img, target)
//  [python]
def test_supported_input_structure(self, T):

        batch_size = 32
        num_classes = 100

        dataset = self.DummyDataset(size=batch_size, num_classes=num_classes)

        cutmix_mixup = T(num_classes=num_classes)

        dl = DataLoader(dataset, batch_size=batch_size)

        # Input sanity checks
        img, target = next(iter(dl))
        input_img_size = img.shape[-3:]
        assert isinstance(img, torch.Tensor) and isinstance(target, torch.Tensor)
        assert target.shape == (batch_size,)

        def check_output(img, target):
            assert img.shape == (batch_size, *input_img_size)
            assert target.shape == (batch_size, num_classes)
            torch.testing.assert_close(target.sum(axis=-1), torch.ones(batch_size))
            num_non_zero_labels = (target != 0).sum(axis=-1)
            assert (num_non_zero_labels == 2).all()

        # After Dataloader, as unpacked input
        img, target = next(iter(dl))
        assert target.shape == (batch_size,)
        img, target = cutmix_mixup(img, target)
        check_output(img, target)

        # After Dataloader, as packed input
        packed_from_dl = next(iter(dl))
        assert isinstance(packed_from_dl, list)
        img, target = cutmix_mixup(packed_from_dl)
        check_output(img, target)

        # As collation function. We expect default_collate to be used by users.
        def collate_fn_1(batch):
            return cutmix_mixup(default_collate(batch))

        def collate_fn_2(batch):
            return cutmix_mixup(*default_collate(batch))

        for collate_fn in (collate_fn_1, collate_fn_2):
            dl = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)
            img, target = next(iter(dl))
            check_output(img, target)
//  [python]
def test_transform_image_correctness(self, quality, color_space, seed):
        image = make_image(color_space=color_space)

        transform = transforms.JPEG(quality=quality)

        with freeze_rng_state():
            torch.manual_seed(seed)
            actual = transform(image)

            torch.manual_seed(seed)
            expected = F.to_image(transform(F.to_pil_image(image)))

        torch.testing.assert_allclose(actual, expected, rtol=0, atol=1)
//  [python]
def test_transform_image_correctness(self, quality, color_space, seed):
        image = make_image(color_space=color_space)

        transform = transforms.JPEG(quality=quality)

        with freeze_rng_state():
            torch.manual_seed(seed)
            actual = transform(image)

            torch.manual_seed(seed)
            expected = F.to_image(transform(F.to_pil_image(image)))

        torch.testing.assert_close(actual, expected, rtol=0, atol=1)
//  [python]
def test__transform_empty_params(self, mocker):
        transform = transforms.RandomIoUCrop(sampler_options=[2.0])
        image = tv_tensors.Image(torch.rand(1, 3, 4, 4))
        bboxes = tv_tensors.BoundingBoxes(torch.tensor([[1, 1, 2, 2]]), format="XYXY", canvas_size=(4, 4))
        label = torch.tensor([1])
        sample = [image, bboxes, label]
        # Let's mock transform._get_params to control the output:
        transform._get_params = mocker.MagicMock(return_value={})
        output = transform(sample)
        torch.testing.assert_allclose(output, sample)
//  [python]
def test__transform_empty_params(self, mocker):
        transform = transforms.RandomIoUCrop(sampler_options=[2.0])
        image = tv_tensors.Image(torch.rand(1, 3, 4, 4))
        bboxes = tv_tensors.BoundingBoxes(torch.tensor([[1, 1, 2, 2]]), format="XYXY", canvas_size=(4, 4))
        label = torch.tensor([1])
        sample = [image, bboxes, label]
        # Let's mock transform._get_params to control the output:
        transform._get_params = mocker.MagicMock(return_value={})
        output = transform(sample)
        torch.testing.assert_close(output, sample)
//  [python]
def _assert_dense_list_equality(
        self, dense_list_1: List[torch.Tensor], dense_list_2: List[torch.Tensor]
    ) -> None:
        self.assertEqual(len(dense_list_1), len(dense_list_2))
        for tensor_1, tensor_2 in zip(dense_list_1, dense_list_2):
            self.assertEqual(tensor_1, tensor_2)
//  [python]
def _assert_dense_list_equality(
        self, dense_list_1: List[torch.Tensor], dense_list_2: List[torch.Tensor]
    ) -> None:
        self.assertEqual(len(dense_list_1), len(dense_list_2))
        for tensor_1, tensor_2 in zip(dense_list_1, dense_list_2):
            torch.testing.assert_close(tensor_1, tensor_2, rtol=0, atol=0)
//  [python]
def test_shared_tables(
        self,
        device: torch.device,
    ) -> None:
        ebc_config = EmbeddingBagConfig(
            name="t1", embedding_dim=4, num_embeddings=10, feature_names=["f1", "f2"]
        )
        ebc = FusedEmbeddingBagCollection(
            tables=[ebc_config],
            optimizer_type=torch.optim.SGD,
            optimizer_kwargs={"lr": 0.02},
            device=device,
        )

        #     0       1        2  <-- batch
        # f1   [0,1] None    [2]
        # f2   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_lengths_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            lengths=torch.tensor([2, 0, 1, 1, 1, 3]),
        ).to(device)

        pooled_embeddings = ebc(features)

        self.assertEqual(pooled_embeddings.keys(), ["f1", "f2"])
        self.assertEqual(pooled_embeddings["f1"].shape, (features.stride(), 4))
        self.assertEqual(pooled_embeddings["f2"].shape, (features.stride(), 4))

        self.assertEqual(pooled_embeddings["f1"][1], torch.zeros(4).to(device))
//  [python]
def test_shared_tables(
        self,
        device: torch.device,
    ) -> None:
        ebc_config = EmbeddingBagConfig(
            name="t1", embedding_dim=4, num_embeddings=10, feature_names=["f1", "f2"]
        )
        ebc = FusedEmbeddingBagCollection(
            tables=[ebc_config],
            optimizer_type=torch.optim.SGD,
            optimizer_kwargs={"lr": 0.02},
            device=device,
        )

        #     0       1        2  <-- batch
        # f1   [0,1] None    [2]
        # f2   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_lengths_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            lengths=torch.tensor([2, 0, 1, 1, 1, 3]),
        ).to(device)

        pooled_embeddings = ebc(features)

        self.assertEqual(pooled_embeddings.keys(), ["f1", "f2"])
        self.assertEqual(pooled_embeddings["f1"].shape, (features.stride(), 4))
        self.assertEqual(pooled_embeddings["f2"].shape, (features.stride(), 4))

        torch.testing.assert_close(
            pooled_embeddings["f1"][1], torch.zeros(4).to(device)
        )
//  [python]
def test_trainer(recipe_name):
    if recipe_name == 'hot' and not torch.cuda.is_available():
        pytest.xfail(
            'SAM currently requires running with mixed precision due to a composer bug.'
        )

    with open('yamls/deeplabv3.yaml') as f:
        base_config = OmegaConf.load(f)

    with open('tests/smoketest_config.yaml') as f:
        smoke_config = OmegaConf.load(f)

    config = OmegaConf.merge(base_config, smoke_config)
    config.recipe_name = recipe_name

    with SynthADE20KDirectory() as tmp_datadir:
        print(tmp_datadir)
        config.train_dataset.path = tmp_datadir
        config.eval_dataset.path = tmp_datadir
        # Also save checkpoints in the temporary directory
        config.save_folder = tmp_datadir

        # Train
        trainer1 = main(config)
        model1 = trainer1.state.model.module

        # Check that the checkpoint was saved
        chkpt_path = os.path.join(tmp_datadir, 'ep0-ba1-rank0.pt')
        assert os.path.isfile(chkpt_path)

        # Check that the checkpoint was loaded by comparing model weights
        config.load_path = chkpt_path
        config.is_train = False
        config.seed += 10  # change seed
        trainer2 = main(config)
        model2 = trainer2.state.model.module

        for param1, param2 in zip(model1.parameters(), model2.parameters()):
            torch.testing.assert_allclose(param1, param2)
//  [python]
def test_trainer(recipe_name):
    if recipe_name == 'hot' and not torch.cuda.is_available():
        pytest.xfail(
            'SAM currently requires running with mixed precision due to a composer bug.'
        )

    with open('yamls/deeplabv3.yaml') as f:
        base_config = OmegaConf.load(f)

    with open('tests/smoketest_config.yaml') as f:
        smoke_config = OmegaConf.load(f)

    config = OmegaConf.merge(base_config, smoke_config)
    config.recipe_name = recipe_name

    with SynthADE20KDirectory() as tmp_datadir:
        print(tmp_datadir)
        config.train_dataset.path = tmp_datadir
        config.eval_dataset.path = tmp_datadir
        # Also save checkpoints in the temporary directory
        config.save_folder = tmp_datadir

        # Train
        trainer1 = main(config)
        model1 = trainer1.state.model.module

        # Check that the checkpoint was saved
        chkpt_path = os.path.join(tmp_datadir, 'ep0-ba1-rank0.pt')
        assert os.path.isfile(chkpt_path)

        # Check that the checkpoint was loaded by comparing model weights
        config.load_path = chkpt_path
        config.is_train = False
        config.seed += 10  # change seed
        trainer2 = main(config)
        model2 = trainer2.state.model.module

        for param1, param2 in zip(model1.parameters(), model2.parameters()):
            torch.testing.assert_close(param1, param2)
//  [python]
def _test_leaky_relu(
        self,
        input_size,
        negative_slope=0.01,
        test_name="leaky_relu",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        slope = Tensor(
            shape=[],
            dtype=dtype,
            name="slope",
            value=negative_slope,
        )
        X2_op = ops.elementwise(FuncEnum.LRELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, slope)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype=dtype)
        OP_pt = torch.nn.LeakyReLU(negative_slope)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_leaky_relu(
        self,
        input_size,
        negative_slope=0.01,
        test_name="leaky_relu",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        slope = Tensor(
            shape=[],
            dtype=dtype,
            name="slope",
            value=negative_slope,
        )
        X2_op = ops.elementwise(FuncEnum.LRELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, slope)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype=dtype)
        OP_pt = torch.nn.LeakyReLU(negative_slope)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)
    target = torch.randn(batch_size, d_hid, device=args.device)
    loss_fn = torch.nn.MSELoss(reduction="sum")

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
    )

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    ScheduleClass = schedule_map[args.schedule]
    schedule = ScheduleClass(stage, args.chunks, loss_fn=loss_fn)

    # Run
    if args.rank == 0:
        schedule.step(x)
    elif args.rank == args.world_size - 1:
        losses = []
        out = schedule.step(target=target, losses=losses)
    else:
        schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x)
        ref_loss = loss_fn(ref_out, target)
        pipe_loss = sum(losses)
        t.assert_close(out, ref_out, rtol=1e-2, atol=5e-3)
        t.assert_close(pipe_loss, ref_loss)
        print(
            f"equivalence test passed pipe_loss={pipe_loss} ref_loss={ref_loss}"
        )
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)
    target = torch.randn(batch_size, d_hid, device=args.device)
    loss_fn = torch.nn.MSELoss(reduction="sum")

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
    )

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    ScheduleClass = schedule_map[args.schedule]
    schedule = ScheduleClass(stage, args.chunks, loss_fn=loss_fn)

    # Run
    if args.rank == 0:
        schedule.step(x)
    elif args.rank == args.world_size - 1:
        losses = []
        out = schedule.step(target=target, losses=losses)
    else:
        schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x)
        ref_loss = loss_fn(ref_out, target)
        pipe_loss = sum(losses)
        torch.testing.assert_close(out, ref_out, rtol=1e-2, atol=5e-3)
        torch.testing.assert_close(pipe_loss, ref_loss)
        print(
            f"equivalence test passed pipe_loss={pipe_loss} ref_loss={ref_loss}"
        )
//  [python]
def _test_floor(
        self,
        input_size,
        test_name="floor",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.FLOOR)

        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = torch.floor(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_floor(
        self,
        input_size,
        test_name="floor",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.FLOOR)

        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = torch.floor(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_floor_div(
        self,
        input_size,
        test_name="floor_div",
        dividend=2,
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        slope = Tensor(
            shape=[],
            dtype=dtype,
            name="input1",
            value=dividend,
        )
        X2_op = ops.elementwise(FuncEnum.FLOOR_DIV)

        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, slope)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = torch.div(x1_pt, dividend, rounding_mode="floor")

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_floor_div(
        self,
        input_size,
        test_name="floor_div",
        dividend=2,
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        slope = Tensor(
            shape=[],
            dtype=dtype,
            name="input1",
            value=dividend,
        )
        X2_op = ops.elementwise(FuncEnum.FLOOR_DIV)

        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, slope)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = torch.div(x1_pt, dividend, rounding_mode="floor")

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_hardtanh(
        self,
        input_size,
        min_val=-1,
        max_val=1,
        test_name="hard_tanh",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_min = Tensor(
            shape=[],
            dtype=dtype,
            name="min_val",
            value=min_val,
            is_input=True,
        )
        X_max = Tensor(
            shape=[],
            dtype=dtype,
            name="max_val",
            value=max_val,
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.HARDTANH)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_min, X_max)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Hardtanh(min_val=min_val, max_val=max_val)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        t.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_hardtanh(
        self,
        input_size,
        min_val=-1,
        max_val=1,
        test_name="hard_tanh",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_min = Tensor(
            shape=[],
            dtype=dtype,
            name="min_val",
            value=min_val,
            is_input=True,
        )
        X_max = Tensor(
            shape=[],
            dtype=dtype,
            name="max_val",
            value=max_val,
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.HARDTANH)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_min, X_max)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Hardtanh(min_val=min_val, max_val=max_val)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_softplus(
        self,
        input_size,
        beta=1.0,
        threshold=20.0,
        test_name="softplus",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_beta = Tensor(
            shape=[],
            dtype=dtype,
            name="beta",
            value=beta,
            is_input=True,
        )
        X_threshold = Tensor(
            shape=[],
            dtype=dtype,
            name="threshold",
            value=threshold,
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.SOFTPLUS)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_beta, X_threshold)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Softplus(beta=beta, threshold=threshold)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_softplus(
        self,
        input_size,
        beta=1.0,
        threshold=20.0,
        test_name="softplus",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_beta = Tensor(
            shape=[],
            dtype=dtype,
            name="beta",
            value=beta,
            is_input=True,
        )
        X_threshold = Tensor(
            shape=[],
            dtype=dtype,
            name="threshold",
            value=threshold,
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.SOFTPLUS)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_beta, X_threshold)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Softplus(beta=beta, threshold=threshold)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_simple_function(
        self, input_size, function, test_name, copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(function)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = TORCH_EQUIVALENTS[function](x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-2, rtol=1e-2, equal_nan=True)
//  [python]
def _test_simple_function(
        self, input_size, function, test_name, copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(function)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        x2_pt = TORCH_EQUIVALENTS[function](x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2, equal_nan=True)
//  [python]
def _test_elu(
        self, input_size, alpha=1.0, test_name="elu", copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_alpha = Tensor(
            shape=[],
            dtype=dtype,
            name="alpha",
            value=alpha,
        )
        X2_op = ops.elementwise(FuncEnum.ELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_alpha)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.ELU(alpha=alpha)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_elu(
        self, input_size, alpha=1.0, test_name="elu", copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_alpha = Tensor(
            shape=[],
            dtype=dtype,
            name="alpha",
            value=alpha,
        )
        X2_op = ops.elementwise(FuncEnum.ELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_alpha)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.ELU(alpha=alpha)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_gelu(self, input_size, test_name="gelu", copy_op=False, dtype="float16"):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.GELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.GELU()
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        t.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_gelu(self, input_size, test_name="gelu", copy_op=False, dtype="float16"):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.GELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.GELU()
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_softsign(
        self,
        input_size,
        test_name="softsign",
        copy_op=False,
        dtype="float16",
    ):
        X1 = Tensor(
            shape=[IntImm(dim) for dim in input_size],
            dtype=dtype,
            name="input",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.SOFTSIGN)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Softsign()
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        t.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_softsign(
        self,
        input_size,
        test_name="softsign",
        copy_op=False,
        dtype="float16",
    ):
        X1 = Tensor(
            shape=[IntImm(dim) for dim in input_size],
            dtype=dtype,
            name="input",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.SOFTSIGN)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")

        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.Softsign()
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_fast_gelu(
        self, input_size, test_name="fast_gelu", copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.FASTGELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.GELU(approximate="tanh")
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        t.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_fast_gelu(
        self, input_size, test_name="fast_gelu", copy_op=False, dtype="float16"
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2_op = ops.elementwise(FuncEnum.FASTGELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.GELU(approximate="tanh")
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_celu(
        self,
        input_size,
        alpha=1.0,
        test_name="celu",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_alpha = Tensor(
            shape=[],
            dtype=dtype,
            name="alpha",
            value=alpha,
        )
        X2_op = ops.elementwise(FuncEnum.CELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_alpha)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.CELU(alpha=alpha)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_celu(
        self,
        input_size,
        alpha=1.0,
        test_name="celu",
        copy_op=False,
        dtype="float16",
    ):
        assert len(input_size) == 2
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X_alpha = Tensor(
            shape=[],
            dtype=dtype,
            name="alpha",
            value=alpha,
        )
        X2_op = ops.elementwise(FuncEnum.CELU)
        if copy_op:
            X2_op = ops.elementwise(**X2_op._get_op_attributes())
        X2 = X2_op(X1, X_alpha)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", f"{test_name}_{dtype}")
        x1_pt = get_random_torch_tensor(input_size, dtype)
        OP_pt = torch.nn.CELU(alpha=alpha)
        x2_pt = OP_pt(x1_pt)

        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_vanilla_replacement(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()
        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))
//  [python]
def test_vanilla_replacement(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()
        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))
//  [python]
def test_module_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.ModuleList([nn.Conv2d(3, 8, kernel_size=3, padding=1) for _ in range(3)])

            def forward(self, x):
                return torch.cat([conv(x) for conv in self.convs], dim=1)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        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, 3 * 8, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_module_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.ModuleList([nn.Conv2d(3, 8, kernel_size=3, padding=1) for _ in range(3)])

            def forward(self, x):
                return torch.cat([conv(x) for conv in self.convs], dim=1)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        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, 3 * 8, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def _test_fast_sigmoid_extreme_value(
        self, atol, rtol, use_tanh_for_sigmoid, testname
    ):
        X1 = Tensor(
            shape=[IntImm(2), IntImm(2)],
            dtype="float16",
            name="input0",
            is_input=True,
        )
        sigmoid_op = ops.elementwise(FuncEnum.SIGMOID)
        X2 = sigmoid_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target(use_tanh_for_sigmoid=use_tanh_for_sigmoid)
        module = compile_model(X2, target, "./tmp", f"fast_sigmoid_{testname}")

        x1_pt = get_random_torch_tensor((2, 2), "float16")
        x1_pt[0, 0] = -6
        x1_pt[0, 1] = -8
        x1_pt[1, 0] = -10
        x1_pt[1, 1] = -12

        x2_pt = torch.sigmoid(x1_pt)
        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_allclose(x2, x2_pt, atol=atol, rtol=rtol, equal_nan=True)
//  [python]
def _test_fast_sigmoid_extreme_value(
        self, atol, rtol, use_tanh_for_sigmoid, testname
    ):
        X1 = Tensor(
            shape=[IntImm(2), IntImm(2)],
            dtype="float16",
            name="input0",
            is_input=True,
        )
        sigmoid_op = ops.elementwise(FuncEnum.SIGMOID)
        X2 = sigmoid_op(X1)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target(use_tanh_for_sigmoid=use_tanh_for_sigmoid)
        module = compile_model(X2, target, "./tmp", f"fast_sigmoid_{testname}")

        x1_pt = get_random_torch_tensor((2, 2), "float16")
        x1_pt[0, 0] = -6
        x1_pt[0, 1] = -8
        x1_pt[1, 0] = -10
        x1_pt[1, 1] = -12

        x2_pt = torch.sigmoid(x1_pt)
        x2 = torch.empty_like(x2_pt)
        module.run_with_tensors([x1_pt], [x2])
        torch.testing.assert_close(x2, x2_pt, atol=atol, rtol=rtol, equal_nan=True)
//  [python]
def test_sequential_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.Sequential(
                    nn.Conv2d(3, 8, kernel_size=3, padding=1),
                    nn.Conv2d(8, 16, kernel_size=3, padding=1),
                )

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

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        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, 16, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_sequential_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.Sequential(
                    nn.Conv2d(3, 8, kernel_size=3, padding=1),
                    nn.Conv2d(8, 16, kernel_size=3, padding=1),
                )

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

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        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))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
//  [python]
def test_static_selective_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 = SkipQuantization(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.quantize_module(module)

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

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_allclose(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_static_selective_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 = SkipQuantization(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.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 assert_close(a, b):
    rtol, atol = {
        torch.float16: (5e-3, 5e-3),
        torch.bfloat16: (3e-2, 2e-2),
        torch.float32: (None, None),
    }[a.dtype]
    t.assert_allclose(a, b, rtol=rtol, atol=atol)
//  [python]
def assert_close(a, b):
    rtol, atol = {
        torch.float16: (5e-3, 5e-3),
        torch.bfloat16: (3e-2, 2e-2),
        torch.float32: (None, None),
    }[a.dtype]
    torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
//  [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)
            self.assertEqual(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))