jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_method_calls_with_optional_arg(self):
        class A(torch.nn.Module):
            def __init__(self):
                super(A, self).__init__()

            # opt arg in script-to-script invocation
            def forward(self, x, two: int = 2):
                return x + two

        class B(torch.nn.Module):
            def __init__(self):
                super(B, self).__init__()
                self.A0 = A()

            # opt arg in Python-to-script invocation
            def forward(self, x, one: int = 1):
                return self.A0(x) + one

        script_module = torch.jit.script(B())
        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter()
        )
        mobile_module = _load_for_lite_interpreter(buffer)

        input = torch.tensor([5])
        script_module_forward_result = script_module.forward(input)
        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_allclose(
            script_module_forward_result,
            mobile_module_forward_result
        )

        # change ref only
        script_module_forward_result = script_module.forward(input, 2)
        self.assertFalse(
            (script_module_forward_result == mobile_module_forward_result)
            .all()
            .item()
        )

        # now both match again
        mobile_module_forward_result = mobile_module.forward(input, 2)
        torch.testing.assert_allclose(
            script_module_forward_result,
            mobile_module_forward_result
        )
//  [python]
def test_method_calls_with_optional_arg(self):
        class A(torch.nn.Module):
            def __init__(self):
                super(A, self).__init__()

            # opt arg in script-to-script invocation
            def forward(self, x, two: int = 2):
                return x + two

        class B(torch.nn.Module):
            def __init__(self):
                super(B, self).__init__()
                self.A0 = A()

            # opt arg in Python-to-script invocation
            def forward(self, x, one: int = 1):
                return self.A0(x) + one

        script_module = torch.jit.script(B())
        buffer = io.BytesIO(
            script_module._save_to_buffer_for_lite_interpreter()
        )
        mobile_module = _load_for_lite_interpreter(buffer)

        input = torch.tensor([5])
        script_module_forward_result = script_module.forward(input)
        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_close(
            script_module_forward_result,
            mobile_module_forward_result
        )

        # change ref only
        script_module_forward_result = script_module.forward(input, 2)
        self.assertFalse(
            (script_module_forward_result == mobile_module_forward_result)
            .all()
            .item()
        )

        # now both match again
        mobile_module_forward_result = mobile_module.forward(input, 2)
        torch.testing.assert_close(
            script_module_forward_result,
            mobile_module_forward_result
        )
//  [python]
def test_pack_unpack_state(self):
        sm = TestScript.DerivedStateModule()
        x = torch.rand(3, 4, dtype=torch.float)
        torch.testing.assert_allclose(sm(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))

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

        torch.testing.assert_allclose(sm.derived, torch.neg(sm.param))

        # Test load paths
        self.assertTrue(imported.unpack_called.item())
        torch.testing.assert_allclose(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
//  [python]
def test_pack_unpack_state(self):
        sm = TestScript.DerivedStateModule()
        x = torch.rand(3, 4, dtype=torch.float)
        torch.testing.assert_close(sm(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))

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

        torch.testing.assert_close(sm.derived, torch.neg(sm.param))

        # Test load paths
        self.assertTrue(imported.unpack_called.item())
        torch.testing.assert_close(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
//  [python]
fused_model = fuse(model)
print(fused_model.code)
inp = torch.randn(5, 1, 1, 1)
torch.testing.assert_allclose(fused_model(inp), model(inp))
//  [python]
fused_model = fuse(model)
print(fused_model.code)
inp = torch.randn(5, 1, 1, 1)
torch.testing.assert_close(fused_model(inp), model(inp))
//  [python]
def test_to_rgb(self):
        cfg = dict(type='SelfSupDataPreprocessor', to_rgb=True)
        processor: SelfSupDataPreprocessor = MODELS.build(cfg)
        self.assertTrue(processor._channel_conversion)

        fake_data = {
            'inputs':
            [torch.randn((2, 3, 224, 224)),
             torch.randn((2, 3, 224, 224))],
            'data_samples': [DataSample(), DataSample()]
        }
        inputs = processor(fake_data)['inputs']
        torch.testing.assert_allclose(fake_data['inputs'][0].flip(1).float(),
                                      inputs[0])
        torch.testing.assert_allclose(fake_data['inputs'][1].flip(1).float(),
                                      inputs[1])
//  [python]
def test_to_rgb(self):
        cfg = dict(type='SelfSupDataPreprocessor', to_rgb=True)
        processor: SelfSupDataPreprocessor = MODELS.build(cfg)
        self.assertTrue(processor._channel_conversion)

        fake_data = {
            'inputs':
            [torch.randn((2, 3, 224, 224)),
             torch.randn((2, 3, 224, 224))],
            'data_samples': [DataSample(), DataSample()]
        }
        inputs = processor(fake_data)['inputs']
        torch.testing.assert_close(fake_data['inputs'][0].flip(1).float(),
                                      inputs[0])
        torch.testing.assert_close(fake_data['inputs'][1].flip(1).float(),
                                      inputs[1])
//  [python]
def test_wrap_constant_batch_size():
    class Tester:
        def __init__(self, batch_size):
            self.batch_size = batch_size

        @wrap_constant_batch_size
        def test_constant_batch_size(self, inputs):
            assert len(inputs) == self.batch_size
            return inputs

    batch_size_test = 8
    for i in range(1, batch_size_test + 1):
        tensor = torch.ones([i, 2, 2])
        out = Tester(batch_size=batch_size_test).test_constant_batch_size(tensor)
        torch.testing.assert_allclose(out, tensor)

    with pytest.raises(ValueError):
        Tester(batch_size=batch_size_test).test_constant_batch_size(
            torch.ones([batch_size_test + 1, 2, 2])
        )
//  [python]
def test_wrap_constant_batch_size():
    class Tester:
        def __init__(self, batch_size):
            self.batch_size = batch_size

        @wrap_constant_batch_size
        def test_constant_batch_size(self, inputs):
            assert len(inputs) == self.batch_size
            return inputs

    batch_size_test = 8
    for i in range(1, batch_size_test + 1):
        tensor = torch.ones([i, 2, 2])
        out = Tester(batch_size=batch_size_test).test_constant_batch_size(tensor)
        torch.testing.assert_close(out, tensor)

    with pytest.raises(ValueError):
        Tester(batch_size=batch_size_test).test_constant_batch_size(
            torch.ones([batch_size_test + 1, 2, 2])
        )
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

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

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

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        torch.testing.assert_allclose(ref_res, res, rtol=1e-2, atol=1e-3)
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

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

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

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        torch.testing.assert_close(ref_res, res, rtol=1e-2, atol=1e-3)
//  [python]
def test_epsilon(eps_start, q_function):
    policy = EpsGreedy(q_function, ExponentialDecay(eps_start))
    for t in range(100):
        assert policy.param() == eps_start
        policy.update()

    policy = EpsGreedy(
        q_function, ExponentialDecay(start=eps_start, end=0.1, decay=100)
    )
    for t in range(100):
        torch.testing.assert_allclose(
            policy.param(), 0.1 + (eps_start - 0.1) * torch.exp(-torch.tensor(t / 100))
        )
        policy.update()
//  [python]
def test_epsilon(eps_start, q_function):
    policy = EpsGreedy(q_function, ExponentialDecay(eps_start))
    for t in range(100):
        assert policy.param() == eps_start
        policy.update()

    policy = EpsGreedy(
        q_function, ExponentialDecay(start=eps_start, end=0.1, decay=100)
    )
    for t in range(100):
        torch.testing.assert_close(
            policy.param(), 0.1 + (eps_start - 0.1) * torch.exp(-torch.tensor(t / 100))
        )
        policy.update()
//  [python]
def test_area_intersection(device):
    boxes_a = Boxes(torch.tensor(
        [(5, 5, 10, 10)] * 14,
        device=device, dtype=torch.float)
    )

    boxes_b = Boxes(torch.tensor([
        # Outside
        (1, 1, 3, 3),
        (11, 11, 13, 13),
        (1, 11, 3, 13),
        (11, 1, 13, 3),

        # Fully inside
        (6, 7, 9, 9),
        (5, 5, 10, 10),

        # Partial overlap by one corner
        (1, 1, 6, 7),
        (9, 9, 13, 13),
        (1, 7, 7, 13),
        (8, 1, 13, 6),

        # Partial overlap by 2 corners
        (1, 1, 13, 7),
        (1, 7, 13, 13),
        (1, 1, 8, 13),
        (8, 1, 13, 13),
    ], device=device, dtype=torch.float))

    expected_areas = torch.tensor([
        0,
        0,
        0,
        0,
        6,
        25,
        2,
        1,
        6,
        2,
        10,
        15,
        15,
        10,
    ], device=device, dtype=torch.float)

    areas = area_intersection(boxes_a, boxes_b)
    torch.testing.assert_allclose(areas, expected_areas)
//  [python]
def test_area_intersection(device):
    boxes_a = Boxes(torch.tensor(
        [(5, 5, 10, 10)] * 14,
        device=device, dtype=torch.float)
    )

    boxes_b = Boxes(torch.tensor([
        # Outside
        (1, 1, 3, 3),
        (11, 11, 13, 13),
        (1, 11, 3, 13),
        (11, 1, 13, 3),

        # Fully inside
        (6, 7, 9, 9),
        (5, 5, 10, 10),

        # Partial overlap by one corner
        (1, 1, 6, 7),
        (9, 9, 13, 13),
        (1, 7, 7, 13),
        (8, 1, 13, 6),

        # Partial overlap by 2 corners
        (1, 1, 13, 7),
        (1, 7, 13, 13),
        (1, 1, 8, 13),
        (8, 1, 13, 13),
    ], device=device, dtype=torch.float))

    expected_areas = torch.tensor([
        0,
        0,
        0,
        0,
        6,
        25,
        2,
        1,
        6,
        2,
        10,
        15,
        15,
        10,
    ], device=device, dtype=torch.float)

    areas = area_intersection(boxes_a, boxes_b)
    torch.testing.assert_close(areas, expected_areas)
//  [python]
def _test_named_params(self, rank, group, config):
        # Get the named parameters before wrapping.
        before_wrap_model = TransformerWithSharedParams(group)
        before_wrap_params = before_wrap_model.named_parameters()

        # Train the model for 1 step.
        model = self.get_wrapped_model(group, cuda_first=False, config=config)
        self._train_for_several_steps(model, 1, autocast=False)

        # Get the named parameters after wrapping to compare.
        after_wrap_params = model.named_parameters()

        if not config["flatten_parameters"]:
            for before_nm, after_nm in zip(before_wrap_params, after_wrap_params):
                assert before_nm[0] == after_nm[0]
        else:
            named_params_flat = [p for p in after_wrap_params][0][0]
            assert "flat_param_0" in named_params_flat

        # Compare name and size under the `summon_full_params` context.
        with model.summon_full_params():
            after_wrap_params = model.named_parameters()

            for before_nm, after_nm_original in zip(before_wrap_params, after_wrap_params):
                assert before_nm[0] == after_nm_original[0]
                torch.testing.assert_allclose(before_nm[1].shape, after_nm_original[1].cpu().shape)
//  [python]
def _test_named_params(self, rank, group, config):
        # Get the named parameters before wrapping.
        before_wrap_model = TransformerWithSharedParams(group)
        before_wrap_params = before_wrap_model.named_parameters()

        # Train the model for 1 step.
        model = self.get_wrapped_model(group, cuda_first=False, config=config)
        self._train_for_several_steps(model, 1, autocast=False)

        # Get the named parameters after wrapping to compare.
        after_wrap_params = model.named_parameters()

        if not config["flatten_parameters"]:
            for before_nm, after_nm in zip(before_wrap_params, after_wrap_params):
                assert before_nm[0] == after_nm[0]
        else:
            named_params_flat = [p for p in after_wrap_params][0][0]
            assert "flat_param_0" in named_params_flat

        # Compare name and size under the `summon_full_params` context.
        with model.summon_full_params():
            after_wrap_params = model.named_parameters()

            for before_nm, after_nm_original in zip(before_wrap_params, after_wrap_params):
                assert before_nm[0] == after_nm_original[0]
                torch.testing.assert_close(before_nm[1].shape, after_nm_original[1].cpu().shape)
//  [python]
def test_align(self):
        sphere_sampling = SphereSampling(0.1, [1, 0, 0])
        sampled = sphere_sampling(self.data)
        torch.testing.assert_allclose(sampled.pos, torch.tensor([[0.0, 0, 0]]))
        self.assertEqual(sampled.labels[0], 0)

        sphere_sampling = SphereSampling(0.1, [1, 0, 0], align_origin=False)
        sampled = sphere_sampling(self.data)
        torch.testing.assert_allclose(sampled.pos, torch.tensor([[1.0, 0, 0]]))
        self.assertEqual(sampled.labels[0], 0)
//  [python]
def test_align(self):
        sphere_sampling = SphereSampling(0.1, [1, 0, 0])
        sampled = sphere_sampling(self.data)
        torch.testing.assert_close(sampled.pos, torch.tensor([[0.0, 0, 0]]))
        self.assertEqual(sampled.labels[0], 0)

        sphere_sampling = SphereSampling(0.1, [1, 0, 0], align_origin=False)
        sampled = sphere_sampling(self.data)
        torch.testing.assert_close(sampled.pos, torch.tensor([[1.0, 0, 0]]))
        self.assertEqual(sampled.labels[0], 0)
//  [python]
def test_forward(self):
        module = LayerNorm2d(num_channels=32, eps=1e-3)
        input_ = torch.rand(1, 28, 28, 32)
        expected_out = F.layer_norm(input_, module.normalized_shape,
                                    module.weight, module.bias, 1e-3)

        torch.testing.assert_allclose(
            module(input_, data_format='channel_last'), expected_out)

        input_ = input_.permute([0, 3, 1, 2])
        expected_out = expected_out.permute([0, 3, 1, 2])
        torch.testing.assert_allclose(
            module(input_, data_format='channel_first'), expected_out)
//  [python]
def test_forward(self):
        module = LayerNorm2d(num_channels=32, eps=1e-3)
        input_ = torch.rand(1, 28, 28, 32)
        expected_out = F.layer_norm(input_, module.normalized_shape,
                                    module.weight, module.bias, 1e-3)

        torch.testing.assert_close(
            module(input_, data_format='channel_last'), expected_out)

        input_ = input_.permute([0, 3, 1, 2])
        expected_out = expected_out.permute([0, 3, 1, 2])
        torch.testing.assert_close(
            module(input_, data_format='channel_first'), expected_out)
//  [python]
def test_multilabel_head(feat):
    head = MultiLabelClsHead()
    fake_gt_label = torch.randint(0, 2, (4, 10))

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, sigmoid=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.sigmoid(logits))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_allclose(features, feat[0])
    else:
        torch.testing.assert_allclose(features, feat)
//  [python]
def test_multilabel_head(feat):
    head = MultiLabelClsHead()
    fake_gt_label = torch.randint(0, 2, (4, 10))

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, sigmoid=False, post_process=False)
    torch.testing.assert_close(pred, torch.sigmoid(logits))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_close(features, feat[0])
    else:
        torch.testing.assert_close(features, feat)
//  [python]
def test_multilabel_linear_head(feat):
    head = MultiLabelLinearClsHead(10, 5)
    fake_gt_label = torch.randint(0, 2, (4, 10))

    head.init_weights()
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, sigmoid=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.sigmoid(logits))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_allclose(features, feat[0])
    else:
        torch.testing.assert_allclose(features, feat)
//  [python]
def test_multilabel_linear_head(feat):
    head = MultiLabelLinearClsHead(10, 5)
    fake_gt_label = torch.randint(0, 2, (4, 10))

    head.init_weights()
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, sigmoid=False, post_process=False)
    torch.testing.assert_close(pred, torch.sigmoid(logits))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_close(features, feat[0])
    else:
        torch.testing.assert_close(features, feat)
//  [python]
def test_linear_head(feat):

    fake_gt_label = torch.randint(0, 10, (4, ))

    # test LinearClsHead forward
    head = LinearClsHead(10, 3)
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test init weights
    head = LinearClsHead(10, 3)
    head.init_weights()
    assert abs(head.fc.weight).sum() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_allclose(features, feat[0])
    else:
        torch.testing.assert_allclose(features, feat)
//  [python]
def test_linear_head(feat):

    fake_gt_label = torch.randint(0, 10, (4, ))

    # test LinearClsHead forward
    head = LinearClsHead(10, 3)
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test init weights
    head = LinearClsHead(10, 3)
    head.init_weights()
    assert abs(head.fc.weight).sum() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_close(features, feat[0])
    else:
        torch.testing.assert_close(features, feat)
//  [python]
def test_cls_head(feat):
    fake_gt_label = torch.randint(0, 10, (4, ))

    # test forward_train with cal_acc=True
    head = ClsHead(cal_acc=True)
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0
    assert 'accuracy' in losses

    # test forward_train with cal_acc=False
    head = ClsHead()
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test forward_train with weight
    weight = torch.tensor([0.5, 0.5, 0.5, 0.5])
    losses_ = head.forward_train(feat, fake_gt_label)
    losses = head.forward_train(feat, fake_gt_label, weight=weight)
    assert losses['loss'].item() == losses_['loss'].item() * 0.5

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_allclose(features, feat[0])
    else:
        torch.testing.assert_allclose(features, feat)
//  [python]
def test_cls_head(feat):
    fake_gt_label = torch.randint(0, 10, (4, ))

    # test forward_train with cal_acc=True
    head = ClsHead(cal_acc=True)
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0
    assert 'accuracy' in losses

    # test forward_train with cal_acc=False
    head = ClsHead()
    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test forward_train with weight
    weight = torch.tensor([0.5, 0.5, 0.5, 0.5])
    losses_ = head.forward_train(feat, fake_gt_label)
    losses = head.forward_train(feat, fake_gt_label, weight=weight)
    assert losses['loss'].item() == losses_['loss'].item() * 0.5

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_close(features, feat[0])
    else:
        torch.testing.assert_close(features, feat)
//  [python]
def test_conformer_head():
    fake_features = ([torch.rand(4, 64), torch.rand(4, 96)], )
    fake_gt_label = torch.randint(0, 10, (4, ))

    # test conformer head forward
    head = ConformerHead(num_classes=10, in_channels=[64, 96])
    losses = head.forward_train(fake_features, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(sum(logits), dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]
//  [python]
def test_conformer_head():
    fake_features = ([torch.rand(4, 64), torch.rand(4, 96)], )
    fake_gt_label = torch.randint(0, 10, (4, ))

    # test conformer head forward
    head = ConformerHead(num_classes=10, in_channels=[64, 96])
    losses = head.forward_train(fake_features, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(sum(logits), dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]
//  [python]
def test_stacked_linear_cls_head(feat):
    # test assertion
    with pytest.raises(AssertionError):
        StackedLinearClsHead(num_classes=3, in_channels=5, mid_channels=10)

    with pytest.raises(AssertionError):
        StackedLinearClsHead(num_classes=-1, in_channels=5, mid_channels=[10])

    fake_gt_label = torch.randint(0, 2, (4, ))  # B, num_classes

    # test forward with default setting
    head = StackedLinearClsHead(
        num_classes=10, in_channels=5, mid_channels=[20])
    head.init_weights()

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    assert features.shape == (4, 20)

    # test forward with full function
    head = StackedLinearClsHead(
        num_classes=3,
        in_channels=5,
        mid_channels=[8, 10],
        dropout_rate=0.2,
        norm_cfg=dict(type='BN1d'),
        act_cfg=dict(type='HSwish'))
    head.init_weights()

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0
//  [python]
def test_stacked_linear_cls_head(feat):
    # test assertion
    with pytest.raises(AssertionError):
        StackedLinearClsHead(num_classes=3, in_channels=5, mid_channels=10)

    with pytest.raises(AssertionError):
        StackedLinearClsHead(num_classes=-1, in_channels=5, mid_channels=[10])

    fake_gt_label = torch.randint(0, 2, (4, ))  # B, num_classes

    # test forward with default setting
    head = StackedLinearClsHead(
        num_classes=10, in_channels=5, mid_channels=[20])
    head.init_weights()

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(feat, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(feat, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(feat)
    assert features.shape == (4, 20)

    # test forward with full function
    head = StackedLinearClsHead(
        num_classes=3,
        in_channels=5,
        mid_channels=[8, 10],
        dropout_rate=0.2,
        norm_cfg=dict(type='BN1d'),
        act_cfg=dict(type='HSwish'))
    head.init_weights()

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0
//  [python]
def test_random_noise(noise_var):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_noise(x, variance=torch.tensor(noise_var, dtype=torch.float))

    torch.testing.assert_allclose(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_random_noise(noise_var):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_noise(x, variance=torch.tensor(noise_var, dtype=torch.float))

    torch.testing.assert_close(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_forward(self):
        print("\nTesting fourier encoding...")
        encoding = FourierEncoding(3, 256)

        self.assertIsNotNone(encoding)

        # Create dummy data
        dummy_input = np.broadcast_to(np.linspace(-1, 1, 3)[None, :], (1, 3))
        dummy_input = torch.tensor(dummy_input, dtype=torch.float32)

        res = encoding(dummy_input)
        torch.testing.assert_allclose(res.shape[-1], 512)

        print("Fourier encoding is executed successfully. \n")
//  [python]
def test_forward(self):
        print("\nTesting fourier encoding...")
        encoding = FourierEncoding(3, 256)

        self.assertIsNotNone(encoding)

        # Create dummy data
        dummy_input = np.broadcast_to(np.linspace(-1, 1, 3)[None, :], (1, 3))
        dummy_input = torch.tensor(dummy_input, dtype=torch.float32)

        res = encoding(dummy_input)
        torch.testing.assert_close(res.shape[-1], 512)

        print("Fourier encoding is executed successfully. \n")
//  [python]
def test_random_brightness(magnitude):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_brightness(x, magnitude=torch.tensor(magnitude, dtype=torch.float))

    torch.testing.assert_allclose(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_random_brightness(magnitude):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_brightness(x, magnitude=torch.tensor(magnitude, dtype=torch.float))

    torch.testing.assert_close(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_csra_head(feat):
    head = CSRAClsHead(num_classes=10, in_channels=20, num_heads=1, lam=0.1)
    fake_gt_label = torch.randint(0, 2, (4, 10))

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_allclose(features, feat[0])
    else:
        torch.testing.assert_allclose(features, feat)
//  [python]
def test_csra_head(feat):
    head = CSRAClsHead(num_classes=10, in_channels=20, num_heads=1, lam=0.1)
    fake_gt_label = torch.randint(0, 2, (4, 10))

    losses = head.forward_train(feat, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(feat)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(feat)
        assert pred.shape == (4, 10)

    # test pre_logits
    features = head.pre_logits(feat)
    if isinstance(feat, tuple):
        torch.testing.assert_close(features, feat[0])
    else:
        torch.testing.assert_close(features, feat)
//  [python]
def test_fold_batch_norms_inference(model, input_shape):
    # Make sure we have non-trivial values to work with
    nn.init.uniform_(model.seq[1].weight)
    nn.init.uniform_(model.seq[1].bias)
    nn.init.uniform_(model.seq[1].running_mean)
    nn.init.uniform_(model.seq[1].running_var)

    model.eval()
    orig_model = deepcopy(model)
    dummy_input = torch.randn(input_shape)
    folded_model = mt.fold_batch_norms(model, dummy_input=dummy_input, inference=True)
    assert type(folded_model.seq[0]) == type(orig_model.seq[0])
    assert type(folded_model.seq[1]) == nn.Identity

    y_orig = orig_model(dummy_input)
    y_folded = folded_model(dummy_input)
    torch.testing.assert_allclose(y_folded, y_orig)
//  [python]
def test_fold_batch_norms_inference(model, input_shape):
    # Make sure we have non-trivial values to work with
    nn.init.uniform_(model.seq[1].weight)
    nn.init.uniform_(model.seq[1].bias)
    nn.init.uniform_(model.seq[1].running_mean)
    nn.init.uniform_(model.seq[1].running_var)

    model.eval()
    orig_model = deepcopy(model)
    dummy_input = torch.randn(input_shape)
    folded_model = mt.fold_batch_norms(model, dummy_input=dummy_input, inference=True)
    assert type(folded_model.seq[0]) == type(orig_model.seq[0])
    assert type(folded_model.seq[1]) == nn.Identity

    y_orig = orig_model(dummy_input)
    y_folded = folded_model(dummy_input)
    torch.testing.assert_close(y_folded, y_orig)
//  [python]
def test_random_contrast(magnitude):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_contrast(x, magnitude=torch.tensor(magnitude, dtype=torch.float))

    torch.testing.assert_allclose(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_random_contrast(magnitude):
    in_channels = 3
    in_height = 224
    in_width = 224
    x = torch.ones(size=(1, in_channels, in_width, in_height), dtype=torch.float)

    aug_out = random_contrast(x, magnitude=torch.tensor(magnitude, dtype=torch.float))

    torch.testing.assert_close(actual=x.shape, expected=aug_out.shape)
//  [python]
def test_efficientformer_head():
    fake_features = (torch.rand(4, 64), )
    fake_gt_label = torch.randint(0, 10, (4, ))

    # Test without distillation head
    head = EfficientFormerClsHead(
        num_classes=10, in_channels=64, distillation=False)

    # test EfficientFormer head forward
    losses = head.forward_train(fake_features, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]

    # Test without distillation head
    head = EfficientFormerClsHead(num_classes=10, in_channels=64)
    assert hasattr(head, 'head')
    assert hasattr(head, 'dist_head')

    # Test loss
    with pytest.raises(NotImplementedError):
        losses = head.forward_train(fake_features, fake_gt_label)

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_allclose(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]
//  [python]
def test_efficientformer_head():
    fake_features = (torch.rand(4, 64), )
    fake_gt_label = torch.randint(0, 10, (4, ))

    # Test without distillation head
    head = EfficientFormerClsHead(
        num_classes=10, in_channels=64, distillation=False)

    # test EfficientFormer head forward
    losses = head.forward_train(fake_features, fake_gt_label)
    assert losses['loss'].item() > 0

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]

    # Test without distillation head
    head = EfficientFormerClsHead(num_classes=10, in_channels=64)
    assert hasattr(head, 'head')
    assert hasattr(head, 'dist_head')

    # Test loss
    with pytest.raises(NotImplementedError):
        losses = head.forward_train(fake_features, fake_gt_label)

    # test simple_test with post_process
    pred = head.simple_test(fake_features)
    assert isinstance(pred, list) and len(pred) == 4
    with patch('torch.onnx.is_in_onnx_export', return_value=True):
        pred = head.simple_test(fake_features)
        assert pred.shape == (4, 10)

    # test simple_test without post_process
    pred = head.simple_test(fake_features, post_process=False)
    assert isinstance(pred, torch.Tensor) and pred.shape == (4, 10)
    logits = head.simple_test(fake_features, softmax=False, post_process=False)
    torch.testing.assert_close(pred, torch.softmax(logits, dim=1))

    # test pre_logits
    features = head.pre_logits(fake_features)
    assert features is fake_features[0]
//  [python]
def test_pad_avails(self) -> None:
        """Test padding availability masks."""
        num_elements = 10
        num_points = 20
        input = torch.ones((num_elements, num_points), dtype=torch.bool)

        elements_padded = pad_avails(input, 20, 0)
        self.assertEqual(elements_padded.dtype, torch.bool)
        self.assertEqual(elements_padded.shape, (20, num_points))
        torch.testing.assert_allclose(elements_padded[10:, :], torch.zeros((10, num_points), dtype=torch.bool))

        points_padded = pad_avails(input, 30, 1)
        self.assertEqual(points_padded.dtype, torch.bool)
        self.assertEqual(points_padded.shape, (num_elements, 30))
        torch.testing.assert_allclose(points_padded[:, 20:], torch.zeros((num_elements, 10), dtype=torch.bool))
//  [python]
def test_pad_avails(self) -> None:
        """Test padding availability masks."""
        num_elements = 10
        num_points = 20
        input = torch.ones((num_elements, num_points), dtype=torch.bool)

        elements_padded = pad_avails(input, 20, 0)
        self.assertEqual(elements_padded.dtype, torch.bool)
        self.assertEqual(elements_padded.shape, (20, num_points))
        torch.testing.assert_close(elements_padded[10:, :], torch.zeros((10, num_points), dtype=torch.bool))

        points_padded = pad_avails(input, 30, 1)
        self.assertEqual(points_padded.dtype, torch.bool)
        self.assertEqual(points_padded.shape, (num_elements, 30))
        torch.testing.assert_close(points_padded[:, 20:], torch.zeros((num_elements, 10), dtype=torch.bool))
//  [python]
def test_normal_cdf(self):
        from scipy.stats import norm

        scheduler = DDPMScheduler(num_train_timesteps=10)
        inferer = ControlNetDiffusionInferer(scheduler=scheduler)
        x = torch.linspace(-10, 10, 20)
        cdf_approx = inferer._approx_standard_normal_cdf(x)
        cdf_true = norm.cdf(x)
        torch.testing.assert_allclose(cdf_approx, cdf_true, atol=1e-3, rtol=1e-5)
//  [python]
def test_normal_cdf(self):
        from scipy.stats import norm

        scheduler = DDPMScheduler(num_train_timesteps=10)
        inferer = ControlNetDiffusionInferer(scheduler=scheduler)
        x = torch.linspace(-10, 10, 20)
        cdf_approx = inferer._approx_standard_normal_cdf(x)
        cdf_true = norm.cdf(x)
        torch.testing.assert_close(cdf_approx, cdf_true, atol=1e-3, rtol=1e-5)
//  [python]
def test_fb_fc_packed(self):
        X = np.random.rand(16, 16).astype(np.float32) - 0.5
        W = np.random.rand(16, 16).astype(np.float32) - 0.5
        b = np.random.rand(16).astype(np.float32) - 0.5

        def fc_op(X, W, b):
            return np.dot(X, W.T) + b

        x_tensor = torch.tensor(X)
        w_tensor = torch.tensor(W)
        b_tensor = torch.tensor(b)
        packed_w_tensor = torch.fbgemm_pack_gemm_matrix_fp16(w_tensor)
        actual_output = torch.fbgemm_linear_fp16_weight(x_tensor, packed_w_tensor, b_tensor)
        expected_output = fc_op(X, W, b)
        torch.testing.assert_allclose(expected_output, actual_output.cpu(), atol=1e-3, rtol=1e-3)
//  [python]
def test_fb_fc_packed(self):
        X = np.random.rand(16, 16).astype(np.float32) - 0.5
        W = np.random.rand(16, 16).astype(np.float32) - 0.5
        b = np.random.rand(16).astype(np.float32) - 0.5

        def fc_op(X, W, b):
            return np.dot(X, W.T) + b

        x_tensor = torch.tensor(X)
        w_tensor = torch.tensor(W)
        b_tensor = torch.tensor(b)
        packed_w_tensor = torch.fbgemm_pack_gemm_matrix_fp16(w_tensor)
        actual_output = torch.fbgemm_linear_fp16_weight(x_tensor, packed_w_tensor, b_tensor)
        expected_output = fc_op(X, W, b)
        torch.testing.assert_close(expected_output, actual_output.cpu(), atol=1e-3, rtol=1e-3)