jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_bias_gelu():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(1024, 1024)
            self.gelu = nn.GELU()

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

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    sch["linear"].decompose()
    sch.trace(flatten=True)
    print(sch.mod.graph)

    def pattern(x, bias):
        x = F.gelu(bias + x)
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 2
    assert subgraph[0][0][1].target == operator.add
    assert subgraph[0][1][1].target == "gelu"
    sch.fuse(subgraph, compiler="TorchScript", name="BiasGeLU")
    assert isinstance(sch["BiasGeLU_0"].mod, torch.jit.ScriptModule)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_fallback_fusion():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = nn.Conv2d(3, 32, 3)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = self.conv(x)
            x = self.relu(x)
            x = x + 1
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    def pattern(x: torch.Tensor):
        x = call_module("conv", x)
        x = F.relu(x)
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 2
    sch.fuse(subgraph, compiler=None, name="FusedConvReLU")
    assert sch.mod.FusedConvReLU_0.__class__.__name__ == "FusedConvReLU"
    assert isinstance(getattr(sch.mod.FusedConvReLU_0, "0"), nn.Conv2d)
    assert isinstance(getattr(sch.mod.FusedConvReLU_0, "1"), nn.ReLU)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 3, 32, 32), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    t.assert_close(out, out_ref)
//  [python]
def test_fallback_fusion():
    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = nn.Conv2d(3, 32, 3)
            self.relu = nn.ReLU()

        def forward(self, x):
            x = self.conv(x)
            x = self.relu(x)
            x = x + 1
            return x

    mod = Model().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))

    def pattern(x: torch.Tensor):
        x = call_module("conv", x)
        x = F.relu(x)
        return x

    subgraph = sch.find(pattern)
    assert len(subgraph[0]) == 2
    sch.fuse(subgraph, compiler=None, name="FusedConvReLU")
    assert sch.mod.FusedConvReLU_0.__class__.__name__ == "FusedConvReLU"
    assert isinstance(getattr(sch.mod.FusedConvReLU_0, "0"), nn.Conv2d)
    assert isinstance(getattr(sch.mod.FusedConvReLU_0, "1"), nn.ReLU)

    sch_model, _ = slapo.build(sch, init_weights=False)
    print(sch_model)

    inp = torch.randn((1, 3, 32, 32), requires_grad=True).cuda()
    out = sch_model(inp)
    out_ref = mod(inp)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_tied_embeddings(self):

        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["tied_embeddings"] = True
        cfg["model"]["tied_softmax"] = False

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(src_vocab, trg_vocab)
        self.assertEqual(model.src_embed, model.trg_embed)
        torch.testing.assert_allclose(
            model.src_embed.lut.weight, model.trg_embed.lut.weight
        )
        self.assertEqual(
            model.src_embed.lut.weight.shape, model.trg_embed.lut.weight.shape
        )
//  [python]
def test_tied_embeddings(self):

        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["tied_embeddings"] = True
        cfg["model"]["tied_softmax"] = False

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(src_vocab, trg_vocab)
        self.assertEqual(model.src_embed, model.trg_embed)
        torch.testing.assert_close(
            model.src_embed.lut.weight, model.trg_embed.lut.weight
        )
        self.assertEqual(
            model.src_embed.lut.weight.shape, model.trg_embed.lut.weight.shape
        )
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")
        out = model(torch.tensor([input_ids]))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]])
        torch.allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,])  # fmt: skip
        torch.allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_model_7b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")
        out = model(torch.tensor([input_ids]))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_model_13b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf", device_map="auto")
        out = model(torch.tensor(input_ids))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]])
        torch.allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273])  # fmt: skip
        torch.allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_model_13b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf", device_map="auto")
        out = model(torch.tensor(input_ids))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_model_13bf_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf", device_map="auto")
        out = model(torch.tensor(input_ids))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]])
        torch.testing.assert_allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513])  # fmt: skip
        torch.testing.assert_allclose(out.mean(-1), EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_13bf_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf", device_map="auto")
        out = model(torch.tensor(input_ids))
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513])  # fmt: skip
        torch.testing.assert_close(out.mean(-1), EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_70b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf", device_map="auto")
        out = model(torch.tensor(input_ids))

        EXPECTED_MEAN = torch.tensor(
            [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]], dtype=torch.float32
        )
        torch.allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        EXPECTED_SLICE = torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312])  # fmt: skip
        torch.allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_model_70b_logits(self):
        input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
        model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf", device_map="auto")
        out = model(torch.tensor(input_ids))

        EXPECTED_MEAN = torch.tensor(
            [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]], dtype=torch.float32
        )
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        EXPECTED_SLICE = torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-5, rtol=1e-5)
//  [python]
def test_tied_softmax(self):

        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["decoder"]["type"] = "transformer"
        cfg["model"]["tied_embeddings"] = False
        cfg["model"]["tied_softmax"] = True
        cfg["model"]["decoder"]["embeddings"]["embedding_dim"] = 64

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(
            model.trg_embed.lut.weight.shape,
            model.decoder.output_layer.weight.shape,
        )

        self.assertTrue(torch.allclose(model.trg_embed.lut.weight, model.decoder.output_layer.weight))
//  [python]
def test_tied_softmax(self):

        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)
        cfg["model"]["decoder"]["type"] = "transformer"
        cfg["model"]["tied_embeddings"] = False
        cfg["model"]["tied_softmax"] = True
        cfg["model"]["decoder"]["embeddings"]["embedding_dim"] = 64

        src_vocab = trg_vocab = self.vocab

        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        self.assertEqual(
            model.trg_embed.lut.weight.shape,
            model.decoder.output_layer.weight.shape,
        )

        torch.testing.assert_close(
            model.trg_embed.lut.weight, model.decoder.output_layer.weight
        )
//  [python]
def test_multi_dim_split_for_one_dim(self) -> None:
        grad = torch.arange(10).reshape(5, 2)
        expected_split_grad = (
            torch.arange(6).reshape(3, 2),
            torch.arange(6, 10).reshape(2, 2),
        )
        t.assert_allclose(
            multi_dim_split(grad, split_size=3), expected_split_grad
        )
//  [python]
def test_multi_dim_split_for_one_dim(self) -> None:
        grad = torch.arange(10).reshape(5, 2)
        expected_split_grad = (
            torch.arange(6).reshape(3, 2),
            torch.arange(6, 10).reshape(2, 2),
        )
        torch.testing.assert_close(
            multi_dim_split(grad, split_size=3), expected_split_grad
        )
//  [python]
def test_multi_dim_split_without_spliting(self) -> None:
        grad = torch.arange(15).reshape(5, 3)
        expected_split_grad = (
            torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14]]),
        )
        torch.testing.assert_allclose(
            multi_dim_split(grad, split_size=5), expected_split_grad
        )
//  [python]
def test_multi_dim_split_without_spliting(self) -> None:
        grad = torch.arange(15).reshape(5, 3)
        expected_split_grad = (
            torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], [12, 13, 14]]),
        )
        torch.testing.assert_close(
            multi_dim_split(grad, split_size=5), expected_split_grad
        )
//  [python]
def test_multi_dim_split_for_two_dim(self) -> None:
        grad = torch.arange(15).reshape(5, 3)
        expected_split_grad = (
            torch.tensor([[0, 1], [3, 4]]),
            torch.tensor([[2], [5]]),
            torch.tensor([[6, 7], [9, 10]]),
            torch.tensor([[8], [11]]),
            torch.tensor([[12, 13]]),
            torch.tensor([[14]]),
        )
        torch.testing.assert_allclose(
            multi_dim_split(grad, split_size=2), expected_split_grad
        )
//  [python]
def test_multi_dim_split_for_two_dim(self) -> None:
        grad = torch.arange(15).reshape(5, 3)
        expected_split_grad = (
            torch.tensor([[0, 1], [3, 4]]),
            torch.tensor([[2], [5]]),
            torch.tensor([[6, 7], [9, 10]]),
            torch.tensor([[8], [11]]),
            torch.tensor([[12, 13]]),
            torch.tensor([[14]]),
        )
        torch.testing.assert_close(
            multi_dim_split(grad, split_size=2), expected_split_grad
        )
//  [python]
def test_tied_src_trg_softmax(self):

        # test source embedding, target embedding, and softmax tying
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)

        cfg["model"]["decoder"]["type"] = "transformer"
        cfg["model"]["tied_embeddings"] = True
        cfg["model"]["tied_softmax"] = True
        cfg["model"]["decoder"]["embeddings"]["embedding_dim"] = 64
        cfg["model"]["encoder"]["embeddings"]["embedding_dim"] = 64

        src_vocab = trg_vocab = self.vocab
        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        src_weight = model.src_embed.lut.weight
        trg_weight = model.trg_embed.lut.weight
        output_weight = model.decoder.output_layer.weight

        torch.testing.assert_allclose(src_weight, trg_weight)
        torch.testing.assert_allclose(src_weight, output_weight)
        self.assertEqual(src_weight.shape, trg_weight.shape)
        self.assertEqual(trg_weight.shape, output_weight.shape)

        output_weight.data.fill_(3.0)
        self.assertEqual(output_weight.sum().item(), 7104)
        self.assertEqual(output_weight.sum().item(), src_weight.sum().item())
        self.assertEqual(output_weight.sum().item(), trg_weight.sum().item())
        self.assertEqual(src_weight.sum().item(), trg_weight.sum().item())
//  [python]
def test_tied_src_trg_softmax(self):

        # test source embedding, target embedding, and softmax tying
        torch.manual_seed(self.seed)
        cfg = copy.deepcopy(self.cfg)

        cfg["model"]["decoder"]["type"] = "transformer"
        cfg["model"]["tied_embeddings"] = True
        cfg["model"]["tied_softmax"] = True
        cfg["model"]["decoder"]["embeddings"]["embedding_dim"] = 64
        cfg["model"]["encoder"]["embeddings"]["embedding_dim"] = 64

        src_vocab = trg_vocab = self.vocab
        model = build_model(cfg["model"], src_vocab=src_vocab, trg_vocab=trg_vocab)

        src_weight = model.src_embed.lut.weight
        trg_weight = model.trg_embed.lut.weight
        output_weight = model.decoder.output_layer.weight

        torch.testing.assert_close(src_weight, trg_weight)
        torch.testing.assert_close(src_weight, output_weight)
        self.assertEqual(src_weight.shape, trg_weight.shape)
        self.assertEqual(trg_weight.shape, output_weight.shape)

        output_weight.data.fill_(3.0)
        self.assertEqual(output_weight.sum().item(), 7104)
        self.assertEqual(output_weight.sum().item(), src_weight.sum().item())
        self.assertEqual(output_weight.sum().item(), trg_weight.sum().item())
        self.assertEqual(src_weight.sum().item(), trg_weight.sum().item())
//  [python]
def test_load_save_without_tied_weights(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.tie_word_embeddings = False
        for model_class in self.all_model_classes:
            model = model_class(config)
            with tempfile.TemporaryDirectory() as d:
                model.save_pretrained(d)

                model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct
                reloaded_state = model_reloaded.state_dict()
                for k, v in model.state_dict().items():
                    self.assertIn(k, reloaded_state, f"Key {k} is missing from reloaded")
                    self.assertTrue(torch.allclose(v, reloaded_state[k]), msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}")
                # Checking there was no complain of missing weights
                self.assertEqual(infos["missing_keys"], [])
//  [python]
def test_load_save_without_tied_weights(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.tie_word_embeddings = False
        for model_class in self.all_model_classes:
            model = model_class(config)
            with tempfile.TemporaryDirectory() as d:
                model.save_pretrained(d)

                model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct
                reloaded_state = model_reloaded.state_dict()
                for k, v in model.state_dict().items():
                    self.assertIn(k, reloaded_state, f"Key {k} is missing from reloaded")
                    torch.testing.assert_close(
                        v, reloaded_state[k], msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}"
                    )
                # Checking there was no complain of missing weights
                self.assertEqual(infos["missing_keys"], [])
//  [python]
def test_async_compile(self):
        async_compile = AsyncCompile()
        compiled_res = async_compile.cuda(_SOURCE_CODE, "so")
        async_compile.wait(globals())

        # Test load and call functions in .so.
        x = torch.rand(5).float().cuda()
        y = torch.rand(5).float().cuda()
        a = 2.0
        expected_y = a * x + y
        res = compiled_res.result().saxpy(
            ctypes.c_int(5),
            ctypes.c_float(a),
            ctypes.c_void_p(x.data_ptr()),
            ctypes.c_void_p(y.data_ptr()),
        )
        self.assertEqual(y, expected_y)
//  [python]
def test_async_compile(self):
        async_compile = AsyncCompile()
        compiled_res = async_compile.cuda(_SOURCE_CODE, "so")
        async_compile.wait(globals())

        # Test load and call functions in .so.
        x = torch.rand(5).float().cuda()
        y = torch.rand(5).float().cuda()
        a = 2.0
        expected_y = a * x + y
        res = compiled_res.result().saxpy(
            ctypes.c_int(5),
            ctypes.c_float(a),
            ctypes.c_void_p(x.data_ptr()),
            ctypes.c_void_p(y.data_ptr()),
        )
        torch.testing.assert_close(y, expected_y)
//  [python]
def test_cuda_load(self):
        # Test both .o and .so compilation.
        object_file_path, object_hash_key, source_code_path0 = CUDACodeCache.compile(
            _SOURCE_CODE, "o"
        )
        dll_wrapper, so_hash_key, source_code_path1 = CUDACodeCache.load(
            _SOURCE_CODE, "so"
        )
        self.assertNotEqual(source_code_path0, source_code_path1)
        self.assertNotEqual(object_hash_key, so_hash_key)

        # Test load and call functions in .so.
        x = torch.rand(10).float().cuda()
        y = torch.rand(10).float().cuda()
        a = 5.0
        expected_y = a * x + y
        res = dll_wrapper.saxpy(
            ctypes.c_int(10),
            ctypes.c_float(a),
            ctypes.c_void_p(x.data_ptr()),
            ctypes.c_void_p(y.data_ptr()),
        )
        self.assertEqual(y, expected_y)
//  [python]
def test_cuda_load(self):
        # Test both .o and .so compilation.
        object_file_path, object_hash_key, source_code_path0 = CUDACodeCache.compile(
            _SOURCE_CODE, "o"
        )
        dll_wrapper, so_hash_key, source_code_path1 = CUDACodeCache.load(
            _SOURCE_CODE, "so"
        )
        self.assertNotEqual(source_code_path0, source_code_path1)
        self.assertNotEqual(object_hash_key, so_hash_key)

        # Test load and call functions in .so.
        x = torch.rand(10).float().cuda()
        y = torch.rand(10).float().cuda()
        a = 5.0
        expected_y = a * x + y
        res = dll_wrapper.saxpy(
            ctypes.c_int(10),
            ctypes.c_float(a),
            ctypes.c_void_p(x.data_ptr()),
            ctypes.c_void_p(y.data_ptr()),
        )
        torch.testing.assert_close(y, expected_y)
//  [python]
def test_multinomial_sampling(self, scores, seeds, offsets, indices, gt):
        output = multinomial_sampling(scores, seeds, offsets, indices)
        t.assert_close(output, gt)
//  [python]
def test_multinomial_sampling(self, scores, seeds, offsets, indices, gt):
        output = multinomial_sampling(scores, seeds, offsets, indices)
        torch.testing.assert_close(output, gt)
//  [python]
def test_can_use_safetensors(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            model_tied = model_class(config)
            with tempfile.TemporaryDirectory() as d:
                try:
                    model_tied.save_pretrained(d, safe_serialization=True)
                except Exception as e:
                    raise Exception(f"Class {model_class.__name__} cannot be saved using safetensors: {e}")

                model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct
                reloaded_state = model_reloaded.state_dict()
                for k, v in model_tied.state_dict().items():
                    self.assertIn(k, reloaded_state, f"Key {k} is missing from reloaded")
                    torch.allclose(
                        v, reloaded_state[k], msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}"
                    )
                # Checking there was no complain of missing weights
                self.assertEqual(infos["missing_keys"], [])

                # Checking the tensor sharing are correct
                ptrs = defaultdict(list)
                for k, v in model_tied.state_dict().items():
                    ptrs[v.data_ptr()].append(k)

                shared_ptrs = {k: v for k, v in ptrs.items() if len(v) > 1}

                for _, shared_names in shared_ptrs.items():
                    reloaded_ptrs = {reloaded_state[k].data_ptr() for k in shared_names}
                    self.assertEqual(
                        len(reloaded_ptrs),
                        1,
                        f"The shared pointers are incorrect, found different pointers for keys {shared_names}",
                    )
//  [python]
def test_can_use_safetensors(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            model_tied = model_class(config)
            with tempfile.TemporaryDirectory() as d:
                try:
                    model_tied.save_pretrained(d, safe_serialization=True)
                except Exception as e:
                    raise Exception(f"Class {model_class.__name__} cannot be saved using safetensors: {e}")

                model_reloaded, infos = model_class.from_pretrained(d, output_loading_info=True)
                # Checking the state dicts are correct
                reloaded_state = model_reloaded.state_dict()
                for k, v in model_tied.state_dict().items():
                    self.assertIn(k, reloaded_state, f"Key {k} is missing from reloaded")
                    torch.testing.assert_close(
                        v, reloaded_state[k], msg=lambda x: f"{model_class.__name__}: Tensor {k}: {x}"
                    )
                # Checking there was no complain of missing weights
                self.assertEqual(infos["missing_keys"], [])

                # Checking the tensor sharing are correct
                ptrs = defaultdict(list)
                for k, v in model_tied.state_dict().items():
                    ptrs[v.data_ptr()].append(k)

                shared_ptrs = {k: v for k, v in ptrs.items() if len(v) > 1}

                for _, shared_names in shared_ptrs.items():
                    reloaded_ptrs = {reloaded_state[k].data_ptr() for k in shared_names}
                    self.assertEqual(
                        len(reloaded_ptrs),
                        1,
                        f"The shared pointers are incorrect, found different pointers for keys {shared_names}",
                    )
//  [python]
def _t5_model(self, is_jit, t5_model, expected_asset_name, test_text):
        """Verify that pre-trained T5 models in torchtext produce
        the same output as the HuggingFace reference implementation.
        """
        expected_asset_path = get_asset_path(expected_asset_name)
        transform = t5_model.transform()
        model = t5_model.get_model()
        model = model.eval()

        if is_jit:
            transform = torch.jit.script(transform)
            model = torch.jit.script(model)

        model_input = transform(test_text)
        if model.encoder_only:
            actual = model(encoder_tokens=model_input)["encoder_output"]
            if not is_jit:
                self._t5_get_encoder(model, model_input, actual)
        else:
            actual = model(encoder_tokens=model_input)["decoder_output"]

        expected = torch.load(expected_asset_path)
        torch.testing.assert_allclose(actual, expected, atol=1e-04, rtol=2.5e-06)
//  [python]
def _t5_model(self, is_jit, t5_model, expected_asset_name, test_text):
        """Verify that pre-trained T5 models in torchtext produce
        the same output as the HuggingFace reference implementation.
        """
        expected_asset_path = get_asset_path(expected_asset_name)
        transform = t5_model.transform()
        model = t5_model.get_model()
        model = model.eval()

        if is_jit:
            transform = torch.jit.script(transform)
            model = torch.jit.script(model)

        model_input = transform(test_text)
        if model.encoder_only:
            actual = model(encoder_tokens=model_input)["encoder_output"]
            if not is_jit:
                self._t5_get_encoder(model, model_input, actual)
        else:
            actual = model(encoder_tokens=model_input)["decoder_output"]

        expected = torch.load(expected_asset_path)
        torch.testing.assert_close(actual, expected, atol=1e-04, rtol=2.5e-06)
//  [python]
def test_load_balancing_loss(self):
        r"""
        Let's make sure we can actually compute the loss and do a backward on it.
        """
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.num_experts = 16
        config.output_router_logits = True
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(config.pad_token_id).to(torch_device)
        model = JambaForCausalLM(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask)
        bs, seqlen = input_ids.shape
        self.assertEqual(result.router_logits[0].shape, (bs * seqlen, config.num_experts))
        self.assertTrue(torch.allclose(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2))

        # First, we make sure that adding padding tokens doesn't change the loss
        # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
        pad_length = 1000
        # Add padding tokens to input_ids
        padding_block = config.pad_token_id * torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(
            torch_device
        )
        padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left
        padded_attention_mask = padded_input_ids.ne(config.pad_token_id).to(torch_device)

        padded_result = model(padded_input_ids, attention_mask=padded_attention_mask)
        self.assertTrue(torch.allclose(result.aux_loss.cpu(), padded_result.aux_loss.cpu(), rtol=1e-4, atol=1e-4))

        # We make sure that the loss of including padding tokens != the loss without padding tokens
        # if attention_mask=None --> we don't exclude padding tokens
        include_padding_result = model(padded_input_ids, attention_mask=None)

        # This is to mimic torch.testing.assert_not_close
        self.assertNotAlmostEqual(include_padding_result.aux_loss.item(), result.aux_loss.item())
//  [python]
def test_load_balancing_loss(self):
        r"""
        Let's make sure we can actually compute the loss and do a backward on it.
        """
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.num_experts = 16
        config.output_router_logits = True
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(config.pad_token_id).to(torch_device)
        model = JambaForCausalLM(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask)
        bs, seqlen = input_ids.shape
        self.assertEqual(result.router_logits[0].shape, (bs * seqlen, config.num_experts))
        torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)

        # First, we make sure that adding padding tokens doesn't change the loss
        # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
        pad_length = 1000
        # Add padding tokens to input_ids
        padding_block = config.pad_token_id * torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(
            torch_device
        )
        padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left
        padded_attention_mask = padded_input_ids.ne(config.pad_token_id).to(torch_device)

        padded_result = model(padded_input_ids, attention_mask=padded_attention_mask)
        torch.testing.assert_close(result.aux_loss.cpu(), padded_result.aux_loss.cpu(), rtol=1e-4, atol=1e-4)

        # We make sure that the loss of including padding tokens != the loss without padding tokens
        # if attention_mask=None --> we don't exclude padding tokens
        include_padding_result = model(padded_input_ids, attention_mask=None)

        # This is to mimic torch.testing.assert_not_close
        self.assertNotAlmostEqual(include_padding_result.aux_loss.item(), result.aux_loss.item())
//  [python]
def test_gpt_checkpointing(dtype, bs, model):
    config = model_configs[model]
    outputs = _test_e2e_checkpointing(bs, dtype, config, checkpoint=False)
    outputs_checkpoint = _test_e2e_checkpointing(bs, dtype, config, checkpoint=True)

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols.update(dict(rtol=2e-2, atol=2e-3))
    for i, (ref, test) in enumerate(zip(outputs, outputs_checkpoint)):
        t.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_gpt_checkpointing(dtype, bs, model):
    config = model_configs[model]
    outputs = _test_e2e_checkpointing(bs, dtype, config, checkpoint=False)
    outputs_checkpoint = _test_e2e_checkpointing(bs, dtype, config, checkpoint=True)

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols.update(dict(rtol=2e-2, atol=2e-3))
    for i, (ref, test) in enumerate(zip(outputs, outputs_checkpoint)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_gpt_fp8_parameters(dtype, bs, model):
    if not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    outputs = _test_gpt_fp8_parameters(bs, dtype, config, False)
    outputs_fp8_params = _test_gpt_fp8_parameters(bs, dtype, config, True)

    # Check that results match
    tols = dict(rtol=0.125, atol=0.0675)
    for i, (ref, test) in enumerate(zip(outputs, outputs_fp8_params)):
        t.assert_allclose(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            rtol=0.125,
            atol=0.0675,
        )
//  [python]
def test_gpt_fp8_parameters(dtype, bs, model):
    if not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    outputs = _test_gpt_fp8_parameters(bs, dtype, config, False)
    outputs_fp8_params = _test_gpt_fp8_parameters(bs, dtype, config, True)

    # Check that results match
    tols = dict(rtol=0.125, atol=0.0675)
    for i, (ref, test) in enumerate(zip(outputs, outputs_fp8_params)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            rtol=0.125,
            atol=0.0675,
        )
//  [python]
def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, fp8_model_params):
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    outputs = _test_e2e_selective_recompute(bs, dtype, config, fp8, fp8_model_params, recompute=False)
    outputs_recompute = _test_e2e_selective_recompute(bs, dtype, config, fp8, fp8_model_params, recompute=True)

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-4
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        t.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, fp8_model_params):
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    outputs = _test_e2e_selective_recompute(bs, dtype, config, fp8, fp8_model_params, recompute=False)
    outputs_recompute = _test_e2e_selective_recompute(bs, dtype, config, fp8, fp8_model_params, recompute=True)

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-4
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_gpt_full_activation_recompute(dtype, bs, model, fp8, fp8_model_params, use_reentrant):
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    if not use_reentrant:
        # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion
        os.environ["NVTE_BIAS_GELU_NVFUSION"] = "0"

    outputs, names = _test_e2e_full_recompute(bs, dtype, config, fp8, fp8_model_params,
                                              recompute=False, use_reentrant=use_reentrant)
    outputs_recompute, _ = _test_e2e_full_recompute(bs, dtype, config, fp8, fp8_model_params,
                                                    recompute=True, use_reentrant=use_reentrant)

    if not use_reentrant:
        # Reset bias+GELU fusion flag to avoid contaminating other tests
        del os.environ["NVTE_BIAS_GELU_NVFUSION"]

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-3
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        torch.testing.assert_allclose(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_gpt_full_activation_recompute(dtype, bs, model, fp8, fp8_model_params, use_reentrant):
    if fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    config = model_configs[model]

    if not use_reentrant:
        # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion
        os.environ["NVTE_BIAS_GELU_NVFUSION"] = "0"

    outputs, names = _test_e2e_full_recompute(bs, dtype, config, fp8, fp8_model_params,
                                              recompute=False, use_reentrant=use_reentrant)
    outputs_recompute, _ = _test_e2e_full_recompute(bs, dtype, config, fp8, fp8_model_params,
                                                    recompute=True, use_reentrant=use_reentrant)

    if not use_reentrant:
        # Reset bias+GELU fusion flag to avoid contaminating other tests
        del os.environ["NVTE_BIAS_GELU_NVFUSION"]

    # Check that results match
    tols = dtype_tols(dtype)
    if dtype in (torch.float16, torch.bfloat16):
        tols["atol"] = 1e-3
    if fp8 or fp8_model_params:
        tols.update(dict(rtol=0.125, atol=0.0675))
    for i, (ref, test) in enumerate(zip(outputs, outputs_recompute)):
        torch.testing.assert_close(
            test,
            ref,
            msg=f"Mismatch in tensor {i}",
            **tols,
        )
//  [python]
def test_benchmark_from_datasets(self):
        d1 = AvalancheDataset(dummy_tensor_dataset())
        d2 = AvalancheDataset(dummy_tensor_dataset())
        d3 = AvalancheDataset(dummy_tensor_dataset())
        d4 = AvalancheDataset(dummy_tensor_dataset())

        bm = benchmark_from_datasets(train=[d1, d2], test=[d3, d4])

        # train stream
        train_s = bm.streams["train"]
        assert len(train_s) == 2
        for eid, (exp, d_orig) in enumerate(zip(train_s, [d1, d2])):
            assert exp.current_experience == eid
            for ii, (x, y) in enumerate(exp.dataset):
                t.assert_close(x, d_orig[ii][0])
                t.assert_close(y, d_orig[ii][1])

        # test stream
        train_s = bm.streams["test"]
        assert len(train_s) == 2
        for eid, (exp, d_orig) in enumerate(zip(train_s, [d3, d4])):
            assert exp.current_experience == eid
            for ii, (x, y) in enumerate(exp.dataset):
                t.assert_close(x, d_orig[ii][0])
                t.assert_close(y, d_orig[ii][1])
//  [python]
def test_benchmark_from_datasets(self):
        d1 = AvalancheDataset(dummy_tensor_dataset())
        d2 = AvalancheDataset(dummy_tensor_dataset())
        d3 = AvalancheDataset(dummy_tensor_dataset())
        d4 = AvalancheDataset(dummy_tensor_dataset())

        bm = benchmark_from_datasets(train=[d1, d2], test=[d3, d4])

        # train stream
        train_s = bm.streams["train"]
        assert len(train_s) == 2
        for eid, (exp, d_orig) in enumerate(zip(train_s, [d1, d2])):
            assert exp.current_experience == eid
            for ii, (x, y) in enumerate(exp.dataset):
                torch.testing.assert_close(x, d_orig[ii][0])
                torch.testing.assert_close(y, d_orig[ii][1])

        # test stream
        train_s = bm.streams["test"]
        assert len(train_s) == 2
        for eid, (exp, d_orig) in enumerate(zip(train_s, [d3, d4])):
            assert exp.current_experience == eid
            for ii, (x, y) in enumerate(exp.dataset):
                torch.testing.assert_close(x, d_orig[ii][0])
                torch.testing.assert_close(y, d_orig[ii][1])
//  [python]
def test_op(M, N, dtype, mode):
    capability = torch.cuda.get_device_capability()
    if capability[0] < 8 and dtype == "bfloat16":
        pytest.skip("Only test bfloat16 on devices with sm >= 80")
    dtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}[dtype]
    # create inputs
    x = torch.randn(M, N, dtype=dtype, device='cuda', requires_grad=True)
    idx = 4 + torch.ones(M, dtype=torch.int64, device='cuda')
    # forward pass
    tt_y = triton.ops.cross_entropy(x, idx)
    th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
    if mode == 'forward':
        torch.testing.assert_allclose(th_y, tt_y)
    # backward pass
    elif mode == 'backward':
        dy = torch.randn_like(tt_y)
        # triton backward
        tt_y.backward(dy)
        tt_dx = x.grad.clone()
        # torch backward
        x.grad = None
        th_y.backward(dy)
        th_dx = x.grad.clone()
        if dtype == torch.float16:
            torch.testing.assert_allclose(th_dx, tt_dx, rtol=0.001, atol=0.001)
        else:
            torch.testing.assert_allclose(th_dx, tt_dx)
//  [python]
def test_op(M, N, dtype, mode):
    capability = torch.cuda.get_device_capability()
    if capability[0] < 8 and dtype == "bfloat16":
        pytest.skip("Only test bfloat16 on devices with sm >= 80")
    dtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}[dtype]
    # create inputs
    x = torch.randn(M, N, dtype=dtype, device='cuda', requires_grad=True)
    idx = 4 + torch.ones(M, dtype=torch.int64, device='cuda')
    # forward pass
    tt_y = triton.ops.cross_entropy(x, idx)
    th_y = torch.nn.CrossEntropyLoss(reduction="none")(x, idx)
    if mode == 'forward':
        torch.testing.assert_close(th_y, tt_y)
    # backward pass
    elif mode == 'backward':
        dy = torch.randn_like(tt_y)
        # triton backward
        tt_y.backward(dy)
        tt_dx = x.grad.clone()
        # torch backward
        x.grad = None
        th_y.backward(dy)
        th_dx = x.grad.clone()
        if dtype == torch.float16:
            torch.testing.assert_close(th_dx, tt_dx, rtol=0.001, atol=0.001)
        else:
            torch.testing.assert_close(th_dx, tt_dx)
//  [python]
def test_label_smoothing(self, tensors, smoothing):
        (input, target_indices, target_onehot) = tensors

        labels_onehot = label_smoothing.smooth_labels(input, target_onehot, smoothing)
        labels_indices = label_smoothing.smooth_labels(input, target_indices, smoothing)
        labels_ref = self.reference_smooth_labels(target_onehot, smoothing)

        self.assertEqual(labels_onehot, labels_ref)
        self.assertEqual(labels_indices, labels_ref)
//  [python]
def test_label_smoothing(self, tensors, smoothing):
        (input, target_indices, target_onehot) = tensors

        labels_onehot = label_smoothing.smooth_labels(input, target_onehot, smoothing)
        labels_indices = label_smoothing.smooth_labels(input, target_indices, smoothing)
        labels_ref = self.reference_smooth_labels(target_onehot, smoothing)

        torch.testing.assert_close(labels_onehot, labels_ref)
        torch.testing.assert_close(labels_indices, labels_ref)
//  [python]
def main(args=None):
    mod = ExampleCode()

    x = torch.randn(batch_size, d_hid)
    y = torch.randn(batch_size, d_hid)
    z = torch.randn(batch_size, d_hid)

    with ArgsChunkSpec((0, 0)), KwargsChunkSpec({"z": 0}):
        pipe = pipeline(
            mod,
            chunks,
            example_args=(x, y),
            example_kwargs={"z": z},
        )

    assert pipe.num_stages == 4

    ref = mod(x, y, z)
    out = pipe(x, y, z)[0]
    torch.testing.assert_allclose(out, ref)
    print(f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref)}")
//  [python]
def main(args=None):
    mod = ExampleCode()

    x = torch.randn(batch_size, d_hid)
    y = torch.randn(batch_size, d_hid)
    z = torch.randn(batch_size, d_hid)

    with ArgsChunkSpec((0, 0)), KwargsChunkSpec({"z": 0}):
        pipe = pipeline(
            mod,
            chunks,
            example_args=(x, y),
            example_kwargs={"z": z},
        )

    assert pipe.num_stages == 4

    ref = mod(x, y, z)
    out = pipe(x, y, z)[0]
    torch.testing.assert_close(out, ref)
    print(f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref)}")
//  [python]
def test_split_dataset_random(self):
        x = torch.rand(32, 10)
        y = torch.arange(
            32
        )  # we use ordered labels to reconstruct the order after shuffling
        dd = AvalancheDataset([TensorDataset(x, y)])

        d1, d2 = split_validation_random(validation_size=0.5, shuffle=True, dataset=dd)
        assert len(d1) + len(d2) == len(dd)

        # check data is shuffled
        iis = []
        for x, ii in d1:
            iis.append(ii)
        for x, ii in d2:
            iis.append(ii)
        assert not all(b >= a for a, b in zip(iis[:-1], iis[1:]))

        # check d1
        for x, ii in d1:
            torch.testing.assert_allclose(x, dd[ii.item()][0])
            torch.testing.assert_allclose(ii, dd[ii.item()][1])
        # check d2
        for x, ii in d2:
            torch.testing.assert_allclose(x, dd[ii.item()][0])
            torch.testing.assert_allclose(ii, dd[ii.item()][1])
//  [python]
def test_split_dataset_random(self):
        x = torch.rand(32, 10)
        y = torch.arange(
            32
        )  # we use ordered labels to reconstruct the order after shuffling
        dd = AvalancheDataset([TensorDataset(x, y)])

        d1, d2 = split_validation_random(validation_size=0.5, shuffle=True, dataset=dd)
        assert len(d1) + len(d2) == len(dd)

        # check data is shuffled
        iis = []
        for x, ii in d1:
            iis.append(ii)
        for x, ii in d2:
            iis.append(ii)
        assert not all(b >= a for a, b in zip(iis[:-1], iis[1:]))

        # check d1
        for x, ii in d1:
            torch.testing.assert_close(x, dd[ii.item()][0])
            torch.testing.assert_close(ii, dd[ii.item()][1])
        # check d2
        for x, ii in d2:
            torch.testing.assert_close(x, dd[ii.item()][0])
            torch.testing.assert_close(ii, dd[ii.item()][1])