jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testUnrollEqualityBackward(self):
        """Check if metagrads match for target/ref net."""
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=diffopt, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                target_meta_loss = target_out[0]

                ref_metagrads = torch.autograd.grad(
                    ref_meta_loss, self.reference_net.parameters()
                )
                target_metagrads = torch.autograd.grad(
                    target_meta_loss, fnet.parameters(time=0)
                )

                # Check that metagrads align
                for rg, tg in zip(ref_metagrads, target_metagrads):
                    torch.testing.assert_close(rg, tg)
//  [python]
def testUnrollEqualityBackwardManualUnroll(self):
        """Check if metagrads match for target/ref net.

        A differentiable optimizer is not used (manual inner loop SGD).
        """
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=None, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                target_meta_loss = target_out[0]

                ref_metagrads = torch.autograd.grad(
                    ref_meta_loss, self.reference_net.parameters()
                )
                target_metagrads = torch.autograd.grad(
                    target_meta_loss, fnet.parameters(time=0)
                )

                # Check that metagrads align
                for rg, tg in zip(ref_metagrads, target_metagrads):
                    torch.testing.assert_allclose(rg, tg)
//  [python]
def testUnrollEqualityBackwardManualUnroll(self):
        """Check if metagrads match for target/ref net.

        A differentiable optimizer is not used (manual inner loop SGD).
        """
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=None, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                target_meta_loss = target_out[0]

                ref_metagrads = torch.autograd.grad(
                    ref_meta_loss, self.reference_net.parameters()
                )
                target_metagrads = torch.autograd.grad(
                    target_meta_loss, fnet.parameters(time=0)
                )

                # Check that metagrads align
                for rg, tg in zip(ref_metagrads, target_metagrads):
                    torch.testing.assert_close(rg, tg)
//  [python]
def testUnrollEqualityForward(self):
        """Check if unrolled patched and reference nets produce same meta loss.
        """
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=diffopt, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                ref_fast_weights = ref_out[1]
                ref_train_losses = ref_out[2]
                ref_train_grads = ref_out[3]
                target_meta_loss = target_out[0]
                target_fast_weights = target_out[1]
                target_train_losses = target_out[2]
                target_train_grads = target_out[3]

                # Check final losses match
                self.assertTrue(
                    torch.allclose(ref_meta_loss, target_meta_loss),
                    msg=(
                        "Ref ({}) and target ({}) metaloss differed on test_it "
                        "{} (mse {})"
                    ).format(
                        ref_meta_loss.item(), target_meta_loss.item(), test_it,
                        (ref_meta_loss - target_meta_loss).pow(2).item()
                    )
                )

                # Check that training losses align
                for rl, tl in zip(ref_train_losses, target_train_losses):
                    torch.testing.assert_allclose(rl, tl)

                # Check that fast weights align
                for rw, tw in zip(ref_fast_weights, target_fast_weights):
                    torch.testing.assert_allclose(rw, tw)

                # Check that grads align
                for rgs, tgs in zip(ref_train_grads, target_train_grads):
                    for rg, tg in zip(rgs, tgs):
                        torch.testing.assert_allclose(rg, tg)
//  [python]
def testUnrollEqualityForward(self):
        """Check if unrolled patched and reference nets produce same meta loss.
        """
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=diffopt, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                ref_fast_weights = ref_out[1]
                ref_train_losses = ref_out[2]
                ref_train_grads = ref_out[3]
                target_meta_loss = target_out[0]
                target_fast_weights = target_out[1]
                target_train_losses = target_out[2]
                target_train_grads = target_out[3]

                # Check final losses match
                self.assertTrue(
                    torch.allclose(ref_meta_loss, target_meta_loss),
                    msg=(
                        "Ref ({}) and target ({}) metaloss differed on test_it "
                        "{} (mse {})"
                    ).format(
                        ref_meta_loss.item(), target_meta_loss.item(), test_it,
                        (ref_meta_loss - target_meta_loss).pow(2).item()
                    )
                )

                # Check that training losses align
                for rl, tl in zip(ref_train_losses, target_train_losses):
                    torch.testing.assert_close(rl, tl)

                # Check that fast weights align
                for rw, tw in zip(ref_fast_weights, target_fast_weights):
                    torch.testing.assert_close(rw, tw)

                # Check that grads align
                for rgs, tgs in zip(ref_train_grads, target_train_grads):
                    for rg, tg in zip(rgs, tgs):
                        torch.testing.assert_close(rg, tg)
//  [python]
def test_superslomo(self):
        # Test extracted from Super-SloMo: https://github.com/avinashpaliwal/Super-SloMo
        # A few interesting things happen here: strided inputs of mixed size,
        # plus outputs of mixed shapes.  The latter characteristic happened to
        # expose a memory corruption bug due to not properly guarding the
        # outputs.
        def eager(t0, t1, t2, t3, t4):
            t5 = torch.mul(t0, t4)
            t6 = torch.mul(t2, t3)
            t7 = torch.mul(t6, t1)
            t9 = torch.add(t5, t7)
            t11 = torch.add(t0, t6)
            ft_p = torch.div(t9, t11)
            return (ft_p, t11, t9, t6)

        t0 = torch.rand(1, 6, 352, 352, device="cuda").transpose(0, 1)
        t1 = torch.rand(6, 3, 352, 352, device="cuda")
        t2 = torch.rand(6, device="cuda")[None, None, None, :].permute(3, 0, 1, 2)
        t3 = torch.rand(6, 1, 352, 352, device="cuda")
        t4 = torch.rand(6, 3, 352, 352, device="cuda")
        inputs = [t0, t1, t2, t3, t4]

        script = torch.jit.script(eager)
        for _ in range(4):
            for pair in zip(script(*inputs), eager(*inputs)):
                test, ref = pair
                torch.testing.assert_allclose(test, ref)
        self.assertAllFused(script.graph_for(*inputs))
//  [python]
def test_superslomo(self):
        # Test extracted from Super-SloMo: https://github.com/avinashpaliwal/Super-SloMo
        # A few interesting things happen here: strided inputs of mixed size,
        # plus outputs of mixed shapes.  The latter characteristic happened to
        # expose a memory corruption bug due to not properly guarding the
        # outputs.
        def eager(t0, t1, t2, t3, t4):
            t5 = torch.mul(t0, t4)
            t6 = torch.mul(t2, t3)
            t7 = torch.mul(t6, t1)
            t9 = torch.add(t5, t7)
            t11 = torch.add(t0, t6)
            ft_p = torch.div(t9, t11)
            return (ft_p, t11, t9, t6)

        t0 = torch.rand(1, 6, 352, 352, device="cuda").transpose(0, 1)
        t1 = torch.rand(6, 3, 352, 352, device="cuda")
        t2 = torch.rand(6, device="cuda")[None, None, None, :].permute(3, 0, 1, 2)
        t3 = torch.rand(6, 1, 352, 352, device="cuda")
        t4 = torch.rand(6, 3, 352, 352, device="cuda")
        inputs = [t0, t1, t2, t3, t4]

        script = torch.jit.script(eager)
        for _ in range(4):
            for pair in zip(script(*inputs), eager(*inputs)):
                test, ref = pair
                torch.testing.assert_close(test, ref)
        self.assertAllFused(script.graph_for(*inputs))
//  [python]
def test_i_time_start_gets_correct_time_one_step_windows():
    # while this directly tests the "mean" result, this is really a test that
    # the data from the correct timestep is piped into the aggregator.
    target_time = 3
    area_weights = torch.ones([4]).to(get_device())
    agg = MeanAggregator(area_weights, target_time=target_time)
    target_data = {"a": torch.zeros([2, 1, 4, 4], device=get_device())}
    for i in range(5):
        sample_data = {
            "a": torch.full([2, 1, 4, 4], fill_value=float(i), device=get_device())
        }
        agg.record_batch(
            1.0,
            target_data=target_data,
            gen_data=sample_data,
            target_data_norm=target_data,
            gen_data_norm=sample_data,
            i_time_start=i,
        )
    logs = agg.get_logs(label="metrics")
    torch.testing.assert_allclose(logs["metrics/weighted_bias/a"], target_time)
//  [python]
def test_i_time_start_gets_correct_time_one_step_windows():
    # while this directly tests the "mean" result, this is really a test that
    # the data from the correct timestep is piped into the aggregator.
    target_time = 3
    area_weights = torch.ones([4]).to(get_device())
    agg = MeanAggregator(area_weights, target_time=target_time)
    target_data = {"a": torch.zeros([2, 1, 4, 4], device=get_device())}
    for i in range(5):
        sample_data = {
            "a": torch.full([2, 1, 4, 4], fill_value=float(i), device=get_device())
        }
        agg.record_batch(
            1.0,
            target_data=target_data,
            gen_data=sample_data,
            target_data_norm=target_data,
            gen_data_norm=sample_data,
            i_time_start=i,
        )
    logs = agg.get_logs(label="metrics")
    torch.testing.assert_close(logs["metrics/weighted_bias/a"], target_time)
//  [python]
def test_unary_ops(self):
        def apply(fn):
            return lambda x: fn(2 * x)

        dtypes = [
            torch.int8,
            torch.uint8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        unary_ops = [
            torch.sigmoid,
            torch.reciprocal,
            torch.neg,
            torch.relu,
            torch.log,
            torch.log10,
            torch.log2,
            torch.exp,
            torch.expm1,
            torch.erf,
            torch.erfc,
            torch.cos,
            torch.sin,
            torch.tan,
            torch.acos,
            torch.asin,
            torch.cosh,
            torch.sinh,
            torch.atan,
            torch.tanh,
            torch.sqrt,
            torch.rsqrt,
            torch.abs,
            torch.ceil,
            torch.floor,
            torch.round,
            torch.trunc,
            torch.frac,
        ]
        devices = ["cuda"]
        for dtype, op, device in product(dtypes, unary_ops, devices):
            try:
                x = self.data_for(dtype, device)
                fn = apply(op)
                ref = fn(x)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x,))
                torch.testing.assert_allclose(ref, t(x))
                self.assertAllFused(t.graph_for(x))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device])
                )
//  [python]
def test_unary_ops(self):
        def apply(fn):
            return lambda x: fn(2 * x)

        dtypes = [
            torch.int8,
            torch.uint8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        unary_ops = [
            torch.sigmoid,
            torch.reciprocal,
            torch.neg,
            torch.relu,
            torch.log,
            torch.log10,
            torch.log2,
            torch.exp,
            torch.expm1,
            torch.erf,
            torch.erfc,
            torch.cos,
            torch.sin,
            torch.tan,
            torch.acos,
            torch.asin,
            torch.cosh,
            torch.sinh,
            torch.atan,
            torch.tanh,
            torch.sqrt,
            torch.rsqrt,
            torch.abs,
            torch.ceil,
            torch.floor,
            torch.round,
            torch.trunc,
            torch.frac,
        ]
        devices = ["cuda"]
        for dtype, op, device in product(dtypes, unary_ops, devices):
            try:
                x = self.data_for(dtype, device)
                fn = apply(op)
                ref = fn(x)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x,))
                torch.testing.assert_close(ref, t(x))
                self.assertAllFused(t.graph_for(x))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device])
                )
//  [python]
def test_speed(self):
        start = time.time()
        a = torch.randn(50000, 3).to(torch.float)
        b = torch.randn(10000, 3).to(torch.float)
        batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
        batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
        R = 1
        samples = 50

        idx, dist = ball_query(
            R,
            samples,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        idx1, dist = ball_query(
            R,
            samples,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        print(time.time() - start)
        torch.testing.assert_allclose(idx1, idx)

        self.assertEqual(idx.shape[0], b.shape[0])
        self.assertEqual(dist.shape[0], b.shape[0])
        self.assertLessEqual(idx.max().item(), len(batch_a))
//  [python]
def test_speed(self):
        start = time.time()
        a = torch.randn(50000, 3).to(torch.float)
        b = torch.randn(10000, 3).to(torch.float)
        batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
        batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
        R = 1
        samples = 50

        idx, dist = ball_query(
            R,
            samples,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        idx1, dist = ball_query(
            R,
            samples,
            a,
            b,
            mode="PARTIAL_DENSE",
            batch_x=batch_a,
            batch_y=batch_b,
            sort=True,
        )
        print(time.time() - start)
        torch.testing.assert_close(idx1, idx)

        self.assertEqual(idx.shape[0], b.shape[0])
        self.assertEqual(dist.shape[0], b.shape[0])
        self.assertLessEqual(idx.max().item(), len(batch_a))
//  [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_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_trivial_graph(self):
        s = torch.full((2, 2), 2)
        tg = torch.jit.script(trivial_graph)
        o_ref = tg(s, s, s)
        tg_a = StaticRuntime(tg)
        o_test = tg_a(s, s, s)[0]
        torch.testing.assert_allclose(o_ref, o_test)
//  [python]
def test_trivial_graph(self):
        s = torch.full((2, 2), 2)
        tg = torch.jit.script(trivial_graph)
        o_ref = tg(s, s, s)
        tg_a = StaticRuntime(tg)
        o_test = tg_a(s, s, s)[0]
        torch.testing.assert_close(o_ref, o_test)
//  [python]
def ort_validate(self, onnx_io, inputs, outputs, tolerate_small_mismatch=False):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
        # compute onnxruntime output prediction
        ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
        ort_outs = ort_session.run(None, ort_inputs)

        for i in range(0, len(outputs)):
            try:
                torch.testing.assert_allclose(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
            except AssertionError as error:
                if tolerate_small_mismatch:
                    self.assertIn("(0.00%)", str(error), str(error))
                else:
                    raise
//  [python]
def ort_validate(self, onnx_io, inputs, outputs, tolerate_small_mismatch=False):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
        # compute onnxruntime output prediction
        ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
        ort_outs = ort_session.run(None, ort_inputs)

        for i in range(0, len(outputs)):
            try:
                torch.testing.assert_close(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
            except AssertionError as error:
                if tolerate_small_mismatch:
                    self.assertIn("(0.00%)", str(error), str(error))
                else:
                    raise
//  [python]
def test_init(self, spatial_dims, size, expected):
        test_filter = self.filter_class(spatial_dims=spatial_dims, size=size)
        torch.testing.assert_allclose(expected, test_filter.filter)
        self.assertIsInstance(test_filter, torch.nn.Module)
//  [python]
def test_init(self, spatial_dims, size, expected):
        test_filter = self.filter_class(spatial_dims=spatial_dims, size=size)
        torch.testing.assert_close(expected, test_filter.filter)
        self.assertIsInstance(test_filter, torch.nn.Module)
//  [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_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        with torch.no_grad():
            src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

        attention.eval()
        attention = torch.jit.script(attention)
        attention.eval()
        o_ref = attention(src, src, src, src_mask)

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        for a, b in zip(o_ref, o_test):
            torch.testing.assert_allclose(a, b)
//  [python]
def test_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        with torch.no_grad():
            src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

        attention.eval()
        attention = torch.jit.script(attention)
        attention.eval()
        o_ref = attention(src, src, src, src_mask)

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        for a, b in zip(o_ref, o_test):
            torch.testing.assert_close(a, b)
//  [python]
def test_scale_coords(self) -> None:
        """
        Test vector feature coordinate scaling.
        """
        scale_value: npt.NDArray[np.float32] = np.array([-2.0, 0.0, -1.0])

        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [2.0, 0.0], [-2.0, 0.0]],
                [[-2.0, 0.0], [2.0, 0.0], [-2.0, 0.0]],
            ]
        )

        result = scale_coords(self.coords, scale_value)
        np.testing.assert_allclose(expected_result, result)

        result = scale_coords(torch.from_numpy(self.coords), torch.from_numpy(scale_value))
        torch.testing.assert_allclose(torch.from_numpy(expected_result), result)
//  [python]
def test_scale_coords(self) -> None:
        """
        Test vector feature coordinate scaling.
        """
        scale_value: npt.NDArray[np.float32] = np.array([-2.0, 0.0, -1.0])

        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [2.0, 0.0], [-2.0, 0.0]],
                [[-2.0, 0.0], [2.0, 0.0], [-2.0, 0.0]],
            ]
        )

        result = scale_coords(self.coords, scale_value)
        np.testing.assert_allclose(expected_result, result)

        result = scale_coords(torch.from_numpy(self.coords), torch.from_numpy(scale_value))
        torch.testing.assert_close(torch.from_numpy(expected_result), result)
//  [python]
def test_xflip_coords(self) -> None:
        """
        Test flipping vector feature coordinates about X-axis.
        """
        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [1.0, 1.0], [-1.0, 1.0]],
                [[-1.0, 0.0], [1.0, -1.0], [-1.0, -1.0]],
            ]
        )

        result = xflip_coords(self.coords)
        np.testing.assert_allclose(expected_result, result)

        result = xflip_coords(torch.from_numpy(self.coords))
        torch.testing.assert_allclose(torch.from_numpy(expected_result), result)
//  [python]
def test_xflip_coords(self) -> None:
        """
        Test flipping vector feature coordinates about X-axis.
        """
        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [1.0, 1.0], [-1.0, 1.0]],
                [[-1.0, 0.0], [1.0, -1.0], [-1.0, -1.0]],
            ]
        )

        result = xflip_coords(self.coords)
        np.testing.assert_allclose(expected_result, result)

        result = xflip_coords(torch.from_numpy(self.coords))
        torch.testing.assert_close(torch.from_numpy(expected_result), result)
//  [python]
def test_yflip_coords(self) -> None:
        """
        Test flipping vector feature coordinates about Y-axis.
        """
        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [-1.0, -1.0], [1.0, -1.0]],
                [[1.0, 0.0], [-1.0, 1.0], [1.0, 1.0]],
            ]
        )

        result = yflip_coords(self.coords)
        np.testing.assert_allclose(expected_result, result)

        result = yflip_coords(torch.from_numpy(self.coords))
        torch.testing.assert_allclose(torch.from_numpy(expected_result), result)
//  [python]
def test_yflip_coords(self) -> None:
        """
        Test flipping vector feature coordinates about Y-axis.
        """
        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[0.0, 0.0], [-1.0, -1.0], [1.0, -1.0]],
                [[1.0, 0.0], [-1.0, 1.0], [1.0, 1.0]],
            ]
        )

        result = yflip_coords(self.coords)
        np.testing.assert_allclose(expected_result, result)

        result = yflip_coords(torch.from_numpy(self.coords))
        torch.testing.assert_close(torch.from_numpy(expected_result), result)
//  [python]
def test_translate_coords(self) -> None:
        """
        Test vector feature coordinate translation.
        """
        translation_value: npt.NDArray[np.float32] = np.array([1.0, 0.0, -1.0])

        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[1.0, 0.0], [0.0, 1.0], [2.0, 1.0]],
                [[2.0, 0.0], [0.0, -1.0], [2.0, -1.0]],
            ]
        )

        result = translate_coords(self.coords, translation_value)
        np.testing.assert_allclose(expected_result, result)

        result = translate_coords(self.coords, translation_value, self.avails)
        expected_result[0][0] = [0.0, 0.0]
        np.testing.assert_allclose(expected_result, result)

        result = translate_coords(
            torch.from_numpy(self.coords), torch.from_numpy(translation_value), torch.from_numpy(self.avails)
        )
        torch.testing.assert_allclose(torch.from_numpy(expected_result), result)
//  [python]
def test_translate_coords(self) -> None:
        """
        Test vector feature coordinate translation.
        """
        translation_value: npt.NDArray[np.float32] = np.array([1.0, 0.0, -1.0])

        expected_result: npt.NDArray[np.float32] = np.array(
            [
                [[1.0, 0.0], [0.0, 1.0], [2.0, 1.0]],
                [[2.0, 0.0], [0.0, -1.0], [2.0, -1.0]],
            ]
        )

        result = translate_coords(self.coords, translation_value)
        np.testing.assert_allclose(expected_result, result)

        result = translate_coords(self.coords, translation_value, self.avails)
        expected_result[0][0] = [0.0, 0.0]
        np.testing.assert_allclose(expected_result, result)

        result = translate_coords(
            torch.from_numpy(self.coords), torch.from_numpy(translation_value), torch.from_numpy(self.avails)
        )
        torch.testing.assert_close(torch.from_numpy(expected_result), result)
//  [python]
def test_mlp(self):
        # Arguments taken from benchmark script, ./bench/dlrm_s_benchmark.sh
        ln_bot = [512, 512, 64]
        sigmoid_bot = -1
        ln_top = [100, 1024, 1024, 1024, 1]
        sigmoid_top = 3
        bot_l = create_mlp(ln_bot, sigmoid_bot)
        bot_l_acc = StaticRuntime(bot_l)
        top_l = create_mlp(ln_top, sigmoid_top)
        top_l_acc = StaticRuntime(top_l)
        with torch.no_grad():
            bot_inp = torch.randn(2048, 512)  # torch.Size([2048, 512])
            top_inp = torch.randn(2048, 100)  # torch.Size([2048, 100])
        ref_bot = bot_l(bot_inp)
        acc_bot = bot_l_acc(bot_inp)[0]
        torch.testing.assert_allclose(acc_bot, ref_bot)
        ref_top = top_l(top_inp)
        acc_top = top_l_acc(top_inp)[0]
        torch.testing.assert_allclose(acc_top, ref_top)
        for _ in range(5):
            with torch.no_grad():
                bot_inp = torch.randn(2048, 512)  # torch.Size([2048, 512])
                top_inp = torch.randn(2048, 100)  # torch.Size([2048, 100])
            ref_bot = bot_l(bot_inp)
            acc_bot = bot_l_acc(bot_inp)[0]
            torch.testing.assert_allclose(acc_bot, ref_bot)
            ref_top = top_l(top_inp)
            acc_top = top_l_acc(top_inp)[0]
            torch.testing.assert_allclose(acc_top, ref_top)
//  [python]
def test_mlp(self):
        # Arguments taken from benchmark script, ./bench/dlrm_s_benchmark.sh
        ln_bot = [512, 512, 64]
        sigmoid_bot = -1
        ln_top = [100, 1024, 1024, 1024, 1]
        sigmoid_top = 3
        bot_l = create_mlp(ln_bot, sigmoid_bot)
        bot_l_acc = StaticRuntime(bot_l)
        top_l = create_mlp(ln_top, sigmoid_top)
        top_l_acc = StaticRuntime(top_l)
        with torch.no_grad():
            bot_inp = torch.randn(2048, 512)  # torch.Size([2048, 512])
            top_inp = torch.randn(2048, 100)  # torch.Size([2048, 100])
        ref_bot = bot_l(bot_inp)
        acc_bot = bot_l_acc(bot_inp)[0]
        torch.testing.assert_close(acc_bot, ref_bot)
        ref_top = top_l(top_inp)
        acc_top = top_l_acc(top_inp)[0]
        torch.testing.assert_close(acc_top, ref_top)
        for _ in range(5):
            with torch.no_grad():
                bot_inp = torch.randn(2048, 512)  # torch.Size([2048, 512])
                top_inp = torch.randn(2048, 100)  # torch.Size([2048, 100])
            ref_bot = bot_l(bot_inp)
            acc_bot = bot_l_acc(bot_inp)[0]
            torch.testing.assert_close(acc_bot, ref_bot)
            ref_top = top_l(top_inp)
            acc_top = top_l_acc(top_inp)[0]
            torch.testing.assert_close(acc_top, ref_top)
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        data_shape = [batch_size] + list(data_shape)
        input_data = torch.rand(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand((weight_output_dim))
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_allclose(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        data_shape = [batch_size] + list(data_shape)
        input_data = torch.rand(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand((weight_output_dim))
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_close(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_conv_gauge_equivariance():
    n_rings = 2
    num_v = 10
    max_order = 4
    gem_precomp = GemPrecomp(n_rings=n_rings, max_order=max_order)
    dtype = torch.float64

    transform_angle = torch.rand(num_v, dtype=dtype) * 2 * pi
    transform = T.Compose(
        (
            SimpleGeometry(),
            gem_precomp,
        )
    )

    transform_t = T.Compose(
        (
            SimpleGeometry(),
            GaugeTransformer(transform_angle),
            gem_precomp,
        )
    )

    data_raw = random_geometry(num_v, edge_p=0.6, dtype=dtype)
    data = transform(data_raw)
    data_t = transform_t(data_raw)

    in_order = 2
    out_order = 2
    channels = (16, 16)
    conv = GemConv(*channels, in_order, out_order, n_rings=n_rings).to(dtype)
    conv.bias.data = torch.rand_like(conv.bias)
    x = torch.randn(num_v, channels[0], 2 * in_order + 1, dtype=dtype)

    x_c = conv(x, data.edge_index, data.precomp, data.connection)
    x_c_t = rep_act(x_c, -transform_angle)

    x_t = rep_act(x, -transform_angle)
    x_t_c = conv(x_t, data_t.edge_index, data_t.precomp, data_t.connection)

    torch.testing.assert_allclose(x_c_t, x_t_c)
//  [python]
def test_conv_gauge_equivariance():
    n_rings = 2
    num_v = 10
    max_order = 4
    gem_precomp = GemPrecomp(n_rings=n_rings, max_order=max_order)
    dtype = torch.float64

    transform_angle = torch.rand(num_v, dtype=dtype) * 2 * pi
    transform = T.Compose(
        (
            SimpleGeometry(),
            gem_precomp,
        )
    )

    transform_t = T.Compose(
        (
            SimpleGeometry(),
            GaugeTransformer(transform_angle),
            gem_precomp,
        )
    )

    data_raw = random_geometry(num_v, edge_p=0.6, dtype=dtype)
    data = transform(data_raw)
    data_t = transform_t(data_raw)

    in_order = 2
    out_order = 2
    channels = (16, 16)
    conv = GemConv(*channels, in_order, out_order, n_rings=n_rings).to(dtype)
    conv.bias.data = torch.rand_like(conv.bias)
    x = torch.randn(num_v, channels[0], 2 * in_order + 1, dtype=dtype)

    x_c = conv(x, data.edge_index, data.precomp, data.connection)
    x_c_t = rep_act(x_c, -transform_angle)

    x_t = rep_act(x, -transform_angle)
    x_t_c = conv(x_t, data_t.edge_index, data_t.precomp, data_t.connection)

    torch.testing.assert_close(x_c_t, x_t_c)
//  [python]
def test_backup_rewards():
    data = arrdict.arrdict(
        v=torch.tensor([[0.], [0.]]),
        w=torch.tensor([[0.], [0.]]),
        n=torch.tensor([0, 0]).int(),
        rewards=torch.tensor([[0.], [1.]]),
        parents=torch.tensor([-1, 0]).int(),
        terminal=torch.tensor([False, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([1]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_allclose(data.w, torch.tensor([[[1.], [1.]]]).cuda())
    torch.testing.assert_allclose(data.n, torch.tensor([[1, 1]]).cuda())
//  [python]
def test_backup_rewards():
    data = arrdict.arrdict(
        v=torch.tensor([[0.], [0.]]),
        w=torch.tensor([[0.], [0.]]),
        n=torch.tensor([0, 0]).int(),
        rewards=torch.tensor([[0.], [1.]]),
        parents=torch.tensor([-1, 0]).int(),
        terminal=torch.tensor([False, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([1]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_close(data.w, torch.tensor([[[1.], [1.]]]).cuda())
    torch.testing.assert_close(data.n, torch.tensor([[1, 1]]).cuda())
//  [python]
def test_root_one_node(device='cuda'):
    data = arrdict.arrdict(
        logits=torch.tensor([[1/3, 2/3]]).log().half(),
        w=torch.tensor([[0.]]).half(),
        n=torch.tensor([0]).short(),
        c_puct=torch.tensor(1.).half(),
        seats=torch.tensor([0]).short(),
        terminal=torch.tensor([False]),
        children=torch.tensor([[-1, -1]]).short())
    
    expected = torch.tensor([[1/3, 2/3]]).half().to(device)
    m = cuda.mcts(**data.to(device)[None])
    actual = cuda.root(m)
    torch.testing.assert_allclose(expected, actual, rtol=1e-3, atol=1e-3)
//  [python]
def test_root_one_node(device='cuda'):
    data = arrdict.arrdict(
        logits=torch.tensor([[1/3, 2/3]]).log().half(),
        w=torch.tensor([[0.]]).half(),
        n=torch.tensor([0]).short(),
        c_puct=torch.tensor(1.).half(),
        seats=torch.tensor([0]).short(),
        terminal=torch.tensor([False]),
        children=torch.tensor([[-1, -1]]).short())
    
    expected = torch.tensor([[1/3, 2/3]]).half().to(device)
    m = cuda.mcts(**data.to(device)[None])
    actual = cuda.root(m)
    torch.testing.assert_close(expected, actual, rtol=1e-3, atol=1e-3)
//  [python]
def test_backup_simple():
    data = arrdict.arrdict(
        v=torch.tensor([[1.], [2.]]),
        w=torch.tensor([[3.], [4.]]),
        n=torch.tensor([5, 6]).int(),
        rewards=torch.tensor([[0.], [0.]]),
        parents=torch.tensor([-1, 0]).int(),
        terminal=torch.tensor([False, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([1]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_allclose(data.w, torch.tensor([[[5.], [6.]]]).cuda())
    torch.testing.assert_allclose(data.n, torch.tensor([[6, 7]]).cuda())
//  [python]
def test_backup_simple():
    data = arrdict.arrdict(
        v=torch.tensor([[1.], [2.]]),
        w=torch.tensor([[3.], [4.]]),
        n=torch.tensor([5, 6]).int(),
        rewards=torch.tensor([[0.], [0.]]),
        parents=torch.tensor([-1, 0]).int(),
        terminal=torch.tensor([False, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([1]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_close(data.w, torch.tensor([[[5.], [6.]]]).cuda())
    torch.testing.assert_close(data.n, torch.tensor([[6, 7]]).cuda())
//  [python]
def test_backup_terminal():
    data = arrdict.arrdict(
        v=torch.tensor([[0.], [1.], [2.]]),
        w=torch.tensor([[0.], [0.], [0.]]),
        n=torch.tensor([0, 0, 0]).int(),
        rewards=torch.tensor([[0.], [3.], [0.]]),
        parents=torch.tensor([-1, 0, 1]).int(),
        terminal=torch.tensor([False, True, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([2]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_allclose(data.w, torch.tensor([[[3.], [3.], [2.]]]).cuda())
//  [python]
def test_backup_terminal():
    data = arrdict.arrdict(
        v=torch.tensor([[0.], [1.], [2.]]),
        w=torch.tensor([[0.], [0.], [0.]]),
        n=torch.tensor([0, 0, 0]).int(),
        rewards=torch.tensor([[0.], [3.], [0.]]),
        parents=torch.tensor([-1, 0, 1]).int(),
        terminal=torch.tensor([False, True, False]))[None].cuda()
    bk = cuda.Backup(**data)

    leaves = torch.tensor([2]).int().cuda()
    cuda.backup(bk, leaves)

    torch.testing.assert_close(data.w, torch.tensor([[[3.], [3.], [2.]]]).cuda())
//  [python]
def test_trivial():
    world = validation.Win.initial(device='cuda')
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=3)

    expected = torch.tensor([[+1.]], device=world.device)
    torch.testing.assert_allclose(m.root().v, expected)
//  [python]
def test_trivial():
    world = validation.Win.initial(device='cuda')
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=3)

    expected = torch.tensor([[+1.]], device=world.device)
    torch.testing.assert_close(m.root().v, expected)