jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_kernel_mutli_gaussian(self, task_name="multi_gauss"):
        keops_kernel = LazyKeopsKernel(
            kernel_type="multi_gauss",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        torch_kernel = TorchKernel(
            kernel_type="multi_gauss",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch.testing.assert_close(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_gaussian_lin(self, task_name="gauss_lin"):
        keops_kernel = LazyKeopsKernel(kernel_type="gauss_lin", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss_lin", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.px, self.py, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.px, self.py, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(),
            (self.x, self.y, self.px, self.py, self.b),
            retain_graph=True,
        )
        torch_grads = grad_torch(
            torch_gauss.mean(),
            (self.x, self.y, self.px, self.py, self.b),
            retain_graph=True,
        )
        torch.testing.assert_allclose(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_gaussian_lin(self, task_name="gauss_lin"):
        keops_kernel = LazyKeopsKernel(kernel_type="gauss_lin", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss_lin", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.px, self.py, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.px, self.py, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(),
            (self.x, self.y, self.px, self.py, self.b),
            retain_graph=True,
        )
        torch_grads = grad_torch(
            torch_gauss.mean(),
            (self.x, self.y, self.px, self.py, self.b),
            retain_graph=True,
        )
        torch.testing.assert_close(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_set_gt_score(self):
        data_sample = ReIDDataSample(metainfo={'num_classes': 5})
        method = getattr(data_sample, 'set_' + 'gt_score')

        # Test set
        score = [0.1, 0.1, 0.6, 0.1, 0.1]
        method(torch.tensor(score))
        sample_gt_label = getattr(data_sample, 'gt_label')
        self.assertIn('score', sample_gt_label)
        torch.testing.assert_allclose(sample_gt_label.score, score)
        self.assertEqual(sample_gt_label.num_classes, 5)

        # Test set again
        score = [0.2, 0.1, 0.5, 0.1, 0.1]
        method(torch.tensor(score))
        torch.testing.assert_allclose(sample_gt_label.score, score)

        # Test invalid type
        with self.assertRaisesRegex(AssertionError, 'be a torch.Tensor'):
            method(score)

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            method(torch.tensor([score]))

        # Test invalid num_classes
        with self.assertRaisesRegex(AssertionError, r'length of value \(6\)'):
            method(torch.tensor(score + [0.1]))

        # Test auto inter num_classes
        data_sample = ReIDDataSample()
        method = getattr(data_sample, 'set_gt_score')
        method(torch.tensor(score))
        sample_gt_label = getattr(data_sample, 'gt_label')
        self.assertEqual(sample_gt_label.num_classes, len(score))
//  [python]
def test_set_gt_score(self):
        data_sample = ReIDDataSample(metainfo={'num_classes': 5})
        method = getattr(data_sample, 'set_' + 'gt_score')

        # Test set
        score = [0.1, 0.1, 0.6, 0.1, 0.1]
        method(torch.tensor(score))
        sample_gt_label = getattr(data_sample, 'gt_label')
        self.assertIn('score', sample_gt_label)
        torch.testing.assert_close(sample_gt_label.score, score)
        self.assertEqual(sample_gt_label.num_classes, 5)

        # Test set again
        score = [0.2, 0.1, 0.5, 0.1, 0.1]
        method(torch.tensor(score))
        torch.testing.assert_close(sample_gt_label.score, score)

        # Test invalid type
        with self.assertRaisesRegex(AssertionError, 'be a torch.Tensor'):
            method(score)

        # Test invalid dims
        with self.assertRaisesRegex(AssertionError, 'but got 2'):
            method(torch.tensor([score]))

        # Test invalid num_classes
        with self.assertRaisesRegex(AssertionError, r'length of value \(6\)'):
            method(torch.tensor(score + [0.1]))

        # Test auto inter num_classes
        data_sample = ReIDDataSample()
        method = getattr(data_sample, 'set_gt_score')
        method(torch.tensor(score))
        sample_gt_label = getattr(data_sample, 'gt_label')
        self.assertEqual(sample_gt_label.num_classes, len(score))
//  [python]
def test_kernel_multi_gaussian_grad(self, task_name="multi_gauss_grad"):
        keops_kernel = LazyKeopsKernel(
            kernel_type="multi_gauss_grad",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        torch_kernel = TorchKernel(
            kernel_type="multi_gauss_grad",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.px, self.x, self.py, self.y)
        torch_gauss = torch_kernel(self.px, self.x, self.py, self.y)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch.testing.assert_allclose(keops_gauss, torch_gauss, rtol=1e-2, atol=1e-5)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_multi_gaussian_grad(self, task_name="multi_gauss_grad"):
        keops_kernel = LazyKeopsKernel(
            kernel_type="multi_gauss_grad",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        torch_kernel = TorchKernel(
            kernel_type="multi_gauss_grad",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.px, self.x, self.py, self.y)
        torch_gauss = torch_kernel(self.px, self.x, self.py, self.y)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch.testing.assert_close(keops_gauss, torch_gauss, rtol=1e-2, atol=1e-5)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [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)
        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)
        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_get_sample(self):
        data = Data(pos=torch.randn((2, 3, 3)))
        torch.testing.assert_allclose(
            MockBaseDataset.get_sample(data, "pos", 1, ConvolutionFormat.DENSE.value), data.pos[1]
        )

        data = Data(pos=torch.randn((3, 3)), batch=torch.tensor([0, 1, 2]))
        torch.testing.assert_allclose(
            MockBaseDataset.get_sample(data, "pos", 1, ConvolutionFormat.PARTIAL_DENSE.value), data.pos[1]
        )
//  [python]
def test_get_sample(self):
        data = Data(pos=torch.randn((2, 3, 3)))
        torch.testing.assert_close(
            MockBaseDataset.get_sample(data, "pos", 1, ConvolutionFormat.DENSE.value), data.pos[1]
        )

        data = Data(pos=torch.randn((3, 3)), batch=torch.tensor([0, 1, 2]))
        torch.testing.assert_close(
            MockBaseDataset.get_sample(data, "pos", 1, ConvolutionFormat.PARTIAL_DENSE.value), data.pos[1]
        )
//  [python]
def test_bug():
    worlds = Hex.initial(n_envs=1, boardsize=3)
    actions = torch.tensor([5, 5, 6, 1], device=worlds.device)
    for a in actions:
        worlds, transitions = worlds.step(a[None])
    torch.testing.assert_allclose(worlds.board[0], torch.tensor([
        [0, 0, 0],
        [5, 0, 1],
        [4, 2, 0]], device=worlds.device))
//  [python]
def test_bug():
    worlds = Hex.initial(n_envs=1, boardsize=3)
    actions = torch.tensor([5, 5, 6, 1], device=worlds.device)
    for a in actions:
        worlds, transitions = worlds.step(a[None])
    torch.testing.assert_close(worlds.board[0], torch.tensor([
        [0, 0, 0],
        [5, 0, 1],
        [4, 2, 0]], device=worlds.device))
//  [python]
def test_inference_long_sentence(self):
        tokenizer = FNetTokenizerFast.from_pretrained("google/fnet-base")

        inputs = tokenizer(
            "the man worked as a [MASK].",
            "this is his [MASK].",
            return_tensors="pt",
            padding="max_length",
            max_length=512,
        )

        # fmt: off
        torch.testing.assert_allclose(inputs["input_ids"], torch.tensor([[4, 13, 283, 2479, 106, 8, 6, 845, 5, 168, 65, 367, 6, 845, 5, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3]]))
        # fmt: on

        inputs = {k: v.to(torch_device) for k, v in inputs.items()}

        model = FNetForMaskedLM.from_pretrained("google/fnet-base")
        model.to(torch_device)
        logits = model(**inputs).logits
        predictions_mask_1 = tokenizer.decode(logits[0, 6].topk(5).indices)
        predictions_mask_2 = tokenizer.decode(logits[0, 12].topk(5).indices)

        self.assertEqual(predictions_mask_1.split(" "), ["man", "child", "teacher", "woman", "model"])
        self.assertEqual(predictions_mask_2.split(" "), ["work", "wife", "job", "story", "name"])
//  [python]
def test_inference_long_sentence(self):
        tokenizer = FNetTokenizerFast.from_pretrained("google/fnet-base")

        inputs = tokenizer(
            "the man worked as a [MASK].",
            "this is his [MASK].",
            return_tensors="pt",
            padding="max_length",
            max_length=512,
        )

        # fmt: off
        torch.testing.assert_close(inputs["input_ids"], torch.tensor([[4, 13, 283, 2479, 106, 8, 6, 845, 5, 168, 65, 367, 6, 845, 5, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3,3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]))
        # fmt: on

        inputs = {k: v.to(torch_device) for k, v in inputs.items()}

        model = FNetForMaskedLM.from_pretrained("google/fnet-base")
        model.to(torch_device)
        logits = model(**inputs).logits
        predictions_mask_1 = tokenizer.decode(logits[0, 6].topk(5).indices)
        predictions_mask_2 = tokenizer.decode(logits[0, 12].topk(5).indices)

        self.assertEqual(predictions_mask_1.split(" "), ["man", "child", "teacher", "woman", "model"])
        self.assertEqual(predictions_mask_2.split(" "), ["work", "wife", "job", "story", "name"])
//  [python]
def test_bug_2():
    worlds = Hex.initial(n_envs=1, boardsize=3)
    worlds.board[:] = torch.tensor([
        [0, 6, 6],
        [1, 1, 1],
        [0, 2, 0]], device=worlds.device, dtype=torch.uint8)
    worlds.seats[:] = 0

    worlds, transitions = worlds.step(torch.tensor([6], device=worlds.device))

    torch.testing.assert_allclose(worlds.board[0], torch.tensor([
        [0, 6, 6],
        [4, 4, 4],
        [4, 2, 0]], device=worlds.device))
//  [python]
def test_bug_2():
    worlds = Hex.initial(n_envs=1, boardsize=3)
    worlds.board[:] = torch.tensor([
        [0, 6, 6],
        [1, 1, 1],
        [0, 2, 0]], device=worlds.device, dtype=torch.uint8)
    worlds.seats[:] = 0

    worlds, transitions = worlds.step(torch.tensor([6], device=worlds.device))

    torch.testing.assert_close(worlds.board[0], torch.tensor([
        [0, 6, 6],
        [4, 4, 4],
        [4, 2, 0]], device=worlds.device))
//  [python]
def test_same_output(self):
        out_cpu = self.forward_cpu(self.input_)
        out_cuda = self.forward_cuda(self.input_)
        torch.testing.assert_allclose(out_cpu, out_cuda.cpu())
//  [python]
def test_same_output(self):
        out_cpu = self.forward_cpu(self.input_)
        out_cuda = self.forward_cuda(self.input_)
        torch.testing.assert_close(out_cpu, out_cuda.cpu())
//  [python]
def _test_equal(batch_dims, num_channels, obj1, obj2, dtype, start, end, tol):
    for dimension in (0, 1, 2):
        sizes = []
        for _ in range(dimension):
            sizes.append(torch.randint(low=1, high=4, size=(1,)).item())
        expected_size = tuple(batch_dims) + tuple(sizes) + (num_channels,)
        eval_times = torch.rand(sizes, dtype=dtype) * (end - start) + start
        obj1_evaluate = obj1.evaluate(eval_times)
        obj2_evaluate = obj2.evaluate(eval_times)
        obj1_derivative = obj1.derivative(eval_times)
        obj2_derivative = obj2.derivative(eval_times)
        assert obj1_evaluate.shape == expected_size
        assert obj2_evaluate.shape == expected_size
        assert obj1_derivative.shape == expected_size
        assert obj2_derivative.shape == expected_size
        torch.testing.assert_allclose(obj1_evaluate, obj2_evaluate, rtol=tol, atol=tol)
        torch.testing.assert_allclose(obj1_derivative, obj2_derivative, rtol=tol, atol=tol)
//  [python]
def _test_equal(batch_dims, num_channels, obj1, obj2, dtype, start, end, tol):
    for dimension in (0, 1, 2):
        sizes = []
        for _ in range(dimension):
            sizes.append(torch.randint(low=1, high=4, size=(1,)).item())
        expected_size = tuple(batch_dims) + tuple(sizes) + (num_channels,)
        eval_times = torch.rand(sizes, dtype=dtype) * (end - start) + start
        obj1_evaluate = obj1.evaluate(eval_times)
        obj2_evaluate = obj2.evaluate(eval_times)
        obj1_derivative = obj1.derivative(eval_times)
        obj2_derivative = obj2.derivative(eval_times)
        assert obj1_evaluate.shape == expected_size
        assert obj2_evaluate.shape == expected_size
        assert obj1_derivative.shape == expected_size
        assert obj2_derivative.shape == expected_size
        torch.testing.assert_close(obj1_evaluate, obj2_evaluate, rtol=tol, atol=tol)
        torch.testing.assert_close(obj1_derivative, obj2_derivative, rtol=tol, atol=tol)
//  [python]
def test_SpatioTemporalClsPositionalEncoding_scriptable(self):
        iter_embed_dim = [1, 2, 4, 32]
        iter_patch_embed_shape = [(1, 1, 1), (1, 2, 4), (32, 16, 1)]
        iter_sep_pos_embed = [True, False]
        iter_has_cls = [True, False]

        for (
            embed_dim,
            patch_embed_shape,
            sep_pos_embed,
            has_cls,
        ) in itertools.product(
            iter_embed_dim,
            iter_patch_embed_shape,
            iter_sep_pos_embed,
            iter_has_cls,
        ):
            stcpe = SpatioTemporalClsPositionalEncoding(
                embed_dim=embed_dim,
                patch_embed_shape=patch_embed_shape,
                sep_pos_embed=sep_pos_embed,
                has_cls=has_cls,
            )
            stcpe_scripted = torch.jit.script(stcpe)
            batch_dim = 4
            video_dim = math.prod(patch_embed_shape)
            fake_input = torch.rand(batch_dim, video_dim, embed_dim)
            expected = stcpe(fake_input)
            actual = stcpe_scripted(fake_input)
            torch.testing.assert_allclose(expected, actual)
//  [python]
def test_SpatioTemporalClsPositionalEncoding_scriptable(self):
        iter_embed_dim = [1, 2, 4, 32]
        iter_patch_embed_shape = [(1, 1, 1), (1, 2, 4), (32, 16, 1)]
        iter_sep_pos_embed = [True, False]
        iter_has_cls = [True, False]

        for (
            embed_dim,
            patch_embed_shape,
            sep_pos_embed,
            has_cls,
        ) in itertools.product(
            iter_embed_dim,
            iter_patch_embed_shape,
            iter_sep_pos_embed,
            iter_has_cls,
        ):
            stcpe = SpatioTemporalClsPositionalEncoding(
                embed_dim=embed_dim,
                patch_embed_shape=patch_embed_shape,
                sep_pos_embed=sep_pos_embed,
                has_cls=has_cls,
            )
            stcpe_scripted = torch.jit.script(stcpe)
            batch_dim = 4
            video_dim = math.prod(patch_embed_shape)
            fake_input = torch.rand(batch_dim, video_dim, embed_dim)
            expected = stcpe(fake_input)
            actual = stcpe_scripted(fake_input)
            torch.testing.assert_close(expected, actual)
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])  # fmt: skip
        torch.testing.assert_allclose(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])  # fmt: skip
        torch.testing.assert_close(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_bilm_batch(self):
        toks = torch.Tensor([[0, 20, 4, 3], [0, 20, 4, 0]]).long()
        lens = torch.Tensor([4, 3]).long()
        idx = pack_padded_sequence(toks, lens, batch_first=True)
        res = self.embedding(idx.data)
        x, xlen = pad_packed_sequence(res)
        tt.assert_allclose(xlen, lens)
        tt.assert_allclose(x, toks)
//  [python]
def test_bilm_batch(self):
        toks = torch.Tensor([[0, 20, 4, 3], [0, 20, 4, 0]]).long()
        lens = torch.Tensor([4, 3]).long()
        idx = pack_padded_sequence(toks, lens, batch_first=True)
        res = self.embedding(idx.data)
        x, xlen = pad_packed_sequence(res)
        tt.assert_close(xlen, lens)
        tt.assert_close(x, toks)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedLayerNorm.bf16_fwd_thresholds
        )
        torch.testing.assert_allclose(actual, expected, **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedLayerNorm.bf16_bwd_thresholds
        )
        torch.testing.assert_allclose(native_x.grad, fused_x.grad, **tols)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedLayerNorm.bf16_fwd_thresholds
        )
        torch.testing.assert_close(actual, expected, **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedLayerNorm.bf16_bwd_thresholds
        )
        torch.testing.assert_close(native_x.grad, fused_x.grad, **tols)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_rms_layers(
            self.normalized_shape, elementwise_affine, dtype
        )
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedRMSNorm.bf16_fwd_thresholds
        )
        torch.testing.assert_allclose(actual, expected.detach().clone().cuda(), **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedRMSNorm.bf16_bwd_thresholds
        )
        torch.testing.assert_allclose(native_x.grad.cuda(), fused_x.grad, **tols)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_rms_layers(
            self.normalized_shape, elementwise_affine, dtype
        )
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedRMSNorm.bf16_fwd_thresholds
        )
        torch.testing.assert_close(actual, expected.detach().clone().cuda(), **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = (
            {"rtol": None, "atol": None}
            if dtype == torch.half
            else TestAutocastFusedRMSNorm.bf16_bwd_thresholds
        )
        torch.testing.assert_close(native_x.grad.cuda(), fused_x.grad, **tols)
//  [python]
def test_loss_dense(self):
        pos = torch.tensor([[[0, 0, 0], [1, 0, 0], [1.1, 0, 0]]], dtype=torch.float)
        f = torch.tensor([[1, 1, 3]], dtype=torch.float)

        var = _variance_estimator_dense(1.01, pos, f)
        torch.testing.assert_allclose(var, [[0, 4, 4]])

        loss = dirichlet_loss(1.01, pos, f)
        self.assertAlmostEqual(loss.item(), 4 / 3.0)
//  [python]
def test_loss_dense(self):
        pos = torch.tensor([[[0, 0, 0], [1, 0, 0], [1.1, 0, 0]]], dtype=torch.float)
        f = torch.tensor([[1, 1, 3]], dtype=torch.float)

        var = _variance_estimator_dense(1.01, pos, f)
        torch.testing.assert_close(var, [[0, 4, 4]])

        loss = dirichlet_loss(1.01, pos, f)
        self.assertAlmostEqual(loss.item(), 4 / 3.0)
//  [python]
def test_prenorm(self, outer_size, inner_size, running_m0, eps, device):
        x = torch.randn(outer_size, inner_size, device=device)
        if running_m0 == 0:
            running_m1 = torch.zeros((inner_size, ), device=device)
            running_m2 = torch.zeros((inner_size, ), device=device)
        else:
            running_m1 = torch.randn((inner_size, ), device=device)
            running_m2 = torch.randn((inner_size, ), device=device)

        m1, m2, scale, bias = self.prenorm(x, running_m0, running_m1,
                                           running_m2, eps)
        m1_ref, m2_ref, scale_ref, bias_ref = prenorm_ref(
            x, running_m0, running_m1, running_m2, eps)

        torch.testing.assert_allclose(m1, m1_ref)
        torch.testing.assert_allclose(m2, m2_ref)
        torch.testing.assert_allclose(scale, scale_ref)
        torch.testing.assert_allclose(bias, bias_ref)
//  [python]
def test_prenorm(self, outer_size, inner_size, running_m0, eps, device):
        x = torch.randn(outer_size, inner_size, device=device)
        if running_m0 == 0:
            running_m1 = torch.zeros((inner_size, ), device=device)
            running_m2 = torch.zeros((inner_size, ), device=device)
        else:
            running_m1 = torch.randn((inner_size, ), device=device)
            running_m2 = torch.randn((inner_size, ), device=device)

        m1, m2, scale, bias = self.prenorm(x, running_m0, running_m1,
                                           running_m2, eps)
        m1_ref, m2_ref, scale_ref, bias_ref = prenorm_ref(
            x, running_m0, running_m1, running_m2, eps)

        torch.testing.assert_close(m1, m1_ref)
        torch.testing.assert_close(m2, m2_ref)
        torch.testing.assert_close(scale, scale_ref)
        torch.testing.assert_close(bias, bias_ref)
//  [python]
def test_compile_api(self):
        # these are mostly checking config processing doesn't blow up with exceptions
        x = torch.randn(8)
        y = dummy_fn(x)
        checks = [
            {},
            {"mode": "default"},
            {"mode": "reduce-overhead"},
            {"mode": "max-autotune"},
            {
                "options": {
                    "max-fusion-size": 128,
                    "unroll_reductions_threshold": 32,
                    "triton.cudagraphs": False,
                }
            },
            {"dynamic": True},
            {"fullgraph": True, "backend": "inductor"},
            {"disable": True},
        ]

        for kwargs in checks:
            torch._dynamo.reset()
            opt_fn = torch.compile(dummy_fn, **kwargs)
            torch.testing.assert_allclose(
                opt_fn(x), y, msg=f"torch.compile(..., **{kwargs!r}) failed"
            )
//  [python]
def test_compile_api(self):
        # these are mostly checking config processing doesn't blow up with exceptions
        x = torch.randn(8)
        y = dummy_fn(x)
        checks = [
            {},
            {"mode": "default"},
            {"mode": "reduce-overhead"},
            {"mode": "max-autotune"},
            {
                "options": {
                    "max-fusion-size": 128,
                    "unroll_reductions_threshold": 32,
                    "triton.cudagraphs": False,
                }
            },
            {"dynamic": True},
            {"fullgraph": True, "backend": "inductor"},
            {"disable": True},
        ]

        for kwargs in checks:
            torch._dynamo.reset()
            opt_fn = torch.compile(dummy_fn, **kwargs)
            torch.testing.assert_close(
                opt_fn(x), y, msg=f"torch.compile(..., **{kwargs!r}) failed"
            )
//  [python]
def test_scripts_properly(self) -> None:
        """
        Test that the VectorMapSimpleMLP model scripts properly.
        """
        model = self._build_model()
        device = torch.device("cpu")
        input_features = self._build_input_features(device, include_agents=True)

        dummy_tensor_input: Dict[str, torch.Tensor] = {}
        dummy_list_tensor_input = {
            "vector_map.coords": input_features["vector_map"].coords,
            "agents.ego": input_features["agents"].ego,
            "agents.agents": input_features["agents"].agents,
        }
        dummy_list_list_tensor_input: Dict[str, List[List[torch.Tensor]]] = {}

        scripted_module = torch.jit.script(model)

        scripted_tensors, scripted_list_tensors, scripted_list_list_tensors = scripted_module.scriptable_forward(
            dummy_tensor_input, dummy_list_tensor_input, dummy_list_list_tensor_input
        )

        py_tensors, py_list_tensors, py_list_list_tensors = model.scriptable_forward(
            dummy_tensor_input, dummy_list_tensor_input, dummy_list_list_tensor_input
        )

        self.assertEqual(1, len(scripted_tensors))
        self.assertEqual(0, len(scripted_list_tensors))
        self.assertEqual(0, len(scripted_list_list_tensors))

        self.assertEqual(1, len(py_tensors))
        self.assertEqual(0, len(py_list_tensors))
        self.assertEqual(0, len(py_list_list_tensors))

        torch.testing.assert_allclose(py_tensors["trajectory"], scripted_tensors["trajectory"])
//  [python]
def test_scripts_properly(self) -> None:
        """
        Test that the VectorMapSimpleMLP model scripts properly.
        """
        model = self._build_model()
        device = torch.device("cpu")
        input_features = self._build_input_features(device, include_agents=True)

        dummy_tensor_input: Dict[str, torch.Tensor] = {}
        dummy_list_tensor_input = {
            "vector_map.coords": input_features["vector_map"].coords,
            "agents.ego": input_features["agents"].ego,
            "agents.agents": input_features["agents"].agents,
        }
        dummy_list_list_tensor_input: Dict[str, List[List[torch.Tensor]]] = {}

        scripted_module = torch.jit.script(model)

        scripted_tensors, scripted_list_tensors, scripted_list_list_tensors = scripted_module.scriptable_forward(
            dummy_tensor_input, dummy_list_tensor_input, dummy_list_list_tensor_input
        )

        py_tensors, py_list_tensors, py_list_list_tensors = model.scriptable_forward(
            dummy_tensor_input, dummy_list_tensor_input, dummy_list_list_tensor_input
        )

        self.assertEqual(1, len(scripted_tensors))
        self.assertEqual(0, len(scripted_list_tensors))
        self.assertEqual(0, len(scripted_list_list_tensors))

        self.assertEqual(1, len(py_tensors))
        self.assertEqual(0, len(py_list_tensors))
        self.assertEqual(0, len(py_list_list_tensors))

        torch.testing.assert_close(py_tensors["trajectory"], scripted_tensors["trajectory"])
//  [python]
def test_loss_sparse(self):
        pos = torch.tensor([[0, 0, 0], [1, 0, 0], [1.1, 0, 0], [0, 0, 0], [1, 0, 0], [1.1, 0, 0]], dtype=torch.float)
        f = torch.tensor([1, 1, 3, 0, 1, 0], dtype=torch.float)
        batch_idx = torch.tensor([0, 0, 0, 1, 1, 1])

        var = _variance_estimator_sparse(1.01, pos, f, batch_idx)
        torch.testing.assert_allclose(var, [0, 4, 4, 1, 2, 1])

        loss = dirichlet_loss(1.01, pos, f, batch_idx)
        self.assertAlmostEqual(loss.item(), sum([0, 4, 4, 1, 2, 1]) / (2 * 6))
//  [python]
def test_loss_sparse(self):
        pos = torch.tensor([[0, 0, 0], [1, 0, 0], [1.1, 0, 0], [0, 0, 0], [1, 0, 0], [1.1, 0, 0]], dtype=torch.float)
        f = torch.tensor([1, 1, 3, 0, 1, 0], dtype=torch.float)
        batch_idx = torch.tensor([0, 0, 0, 1, 1, 1])

        var = _variance_estimator_sparse(1.01, pos, f, batch_idx)
        torch.testing.assert_close(var, [0, 4, 4, 1, 2, 1])

        loss = dirichlet_loss(1.01, pos, f, batch_idx)
        self.assertAlmostEqual(loss.item(), sum([0, 4, 4, 1, 2, 1]) / (2 * 6))
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        # fmt: off
        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])
        # fmt: on
        torch.testing.assert_allclose(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        # fmt: off
        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])
        # fmt: on
        torch.testing.assert_close(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_fwd_thresholds
        torch.testing.assert_allclose(actual, expected, **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_bwd_thresholds
        torch.testing.assert_allclose(native_x.grad, fused_x.grad, **tols)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x)
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_fwd_thresholds
        torch.testing.assert_close(actual, expected, **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.clone()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_bwd_thresholds
        torch.testing.assert_close(native_x.grad, fused_x.grad, **tols)
//  [python]
def test_roll(shifts_dims_pair):
    opset = 11
    from mmcv.onnx.symbolic import register_extra_symbolics
    register_extra_symbolics(opset)

    input = torch.arange(0, 4 * 5 * 6, dtype=torch.float32).view(4, 5, 6)

    shifts, dims = shifts_dims_pair
    func = partial(torch.roll, shifts=shifts, dims=dims)
    wrapped_model = WrapFunction(func).eval()

    with torch.no_grad():
        torch.onnx.export(
            wrapped_model,
            input,
            onnx_file,
            export_params=True,
            keep_initializers_as_inputs=True,
            input_names=['input'],
            output_names=['output'],
            opset_version=opset)

    onnx_model = onnx.load(onnx_file)
    input_all = [node.name for node in onnx_model.graph.input]
    input_initializer = [node.name for node in onnx_model.graph.initializer]
    net_feed_input = list(set(input_all) - set(input_initializer))
    assert (len(net_feed_input) == 1)

    sess = rt.InferenceSession(onnx_file, providers=['CPUExecutionProvider'])
    ort_output = sess.run(None, {'input': input.detach().numpy()})[0]

    with torch.no_grad():
        pytorch_output = wrapped_model(input.clone())

    torch.testing.assert_allclose(ort_output, pytorch_output)
//  [python]
def test_roll(shifts_dims_pair):
    opset = 11
    from mmcv.onnx.symbolic import register_extra_symbolics
    register_extra_symbolics(opset)

    input = torch.arange(0, 4 * 5 * 6, dtype=torch.float32).view(4, 5, 6)

    shifts, dims = shifts_dims_pair
    func = partial(torch.roll, shifts=shifts, dims=dims)
    wrapped_model = WrapFunction(func).eval()

    with torch.no_grad():
        torch.onnx.export(
            wrapped_model,
            input,
            onnx_file,
            export_params=True,
            keep_initializers_as_inputs=True,
            input_names=['input'],
            output_names=['output'],
            opset_version=opset)

    onnx_model = onnx.load(onnx_file)
    input_all = [node.name for node in onnx_model.graph.input]
    input_initializer = [node.name for node in onnx_model.graph.initializer]
    net_feed_input = list(set(input_all) - set(input_initializer))
    assert (len(net_feed_input) == 1)

    sess = rt.InferenceSession(onnx_file, providers=['CPUExecutionProvider'])
    ort_output = sess.run(None, {'input': input.detach().numpy()})[0]

    with torch.no_grad():
        pytorch_output = wrapped_model(input.clone())

    torch.testing.assert_close(ort_output, pytorch_output)
//  [python]
def test_batch_eigh():
#    A = torch.rand(8, 4, 4).cuda()
#    A = A.transpose(1, 2).matmul(A)
    
    num_q=8
    input_lrf = torch.randn(num_q, 8, 4).cuda()
    input_lrf=F.normalize(input_lrf, p=2, dim=-1)
    input_lrf=input_lrf.view(-1,4)
    test=torch.bmm(input_lrf.unsqueeze(2),input_lrf.unsqueeze(1))
    test=test.view(num_q,8,4,4)
    A =torch.sum(test,1)
    A.requires_grad=True
    w, V = torch_cusolver.cusolver_batch_eigh(A,
                                              False,
                                              True,
                                              1e-7,
                                              100,
                                              True)
    (w.sum()+V.sum()).backward()
#    test=torch.symeig(A[0])
    for i in range(A.shape[0]):
        a = A[i]
        e = V[i].t().matmul(w[i].diag()).matmul(V[i])
        torch.testing.assert_allclose(a, e)
//  [python]
def test_batch_eigh():
#    A = torch.rand(8, 4, 4).cuda()
#    A = A.transpose(1, 2).matmul(A)
    
    num_q=8
    input_lrf = torch.randn(num_q, 8, 4).cuda()
    input_lrf=F.normalize(input_lrf, p=2, dim=-1)
    input_lrf=input_lrf.view(-1,4)
    test=torch.bmm(input_lrf.unsqueeze(2),input_lrf.unsqueeze(1))
    test=test.view(num_q,8,4,4)
    A =torch.sum(test,1)
    A.requires_grad=True
    w, V = torch_cusolver.cusolver_batch_eigh(A,
                                              False,
                                              True,
                                              1e-7,
                                              100,
                                              True)
    (w.sum()+V.sum()).backward()
#    test=torch.symeig(A[0])
    for i in range(A.shape[0]):
        a = A[i]
        e = V[i].t().matmul(w[i].diag()).matmul(V[i])
        torch.testing.assert_close(a, e)
//  [python]
def test_load_mobile_module(self):
        class MyTestModule(torch.nn.Module):
            def __init__(self):
                super(MyTestModule, self).__init__()

            def forward(self, x):
                return x + 10

        input = torch.tensor([1])

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

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

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

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

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_allclose(script_module_result, mobile_module_run_method_result)
//  [python]
def test_load_mobile_module(self):
        class MyTestModule(torch.nn.Module):
            def __init__(self):
                super(MyTestModule, self).__init__()

            def forward(self, x):
                return x + 10

        input = torch.tensor([1])

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

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

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

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

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_close(script_module_result, mobile_module_run_method_result)