jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_config_flag_is_respected(self):
        def split_with_cat(x):
            fs = torch.split(x, [4, 4, 24], dim=-1)
            item0 = fs[0]
            item1 = fs[1]
            item2 = fs[2]

            final_items = [item0, item1]
            final_items.extend(item2.split((4, 4, 4, 4, 4, 4), 1))

            return torch.cat(final_items, dim=1)

        args = [
            torch.randn(2, 32),
        ]

        expected = split_with_cat(*args)
        actual = torch.compile(split_with_cat)(*args)

        torch.testing.assert_close(actual, expected)
        self.assertEqual(
            counters["inductor"]["merge_splits_pass"],
            0,
        )
        self.assertEqual(
            counters["inductor"]["normalization_pass"],
            0,
        )
//  [python]
def test_split_cat_merge_mutation(self):
        args = [
            torch.randn(2, 32, 32, 16),
        ]

        def split_cat_mutation(x):
            splits = torch.split(x, 4, dim=1)
            splits[1].copy_(splits[0])
            return torch.cat(splits, dim=1)

        expected = split_cat_mutation(*args)
        actual = torch.compile(split_cat_mutation)(*args)

        self.assertEqual(actual, expected)

        self.assertEqual(counters["inductor"]["scmerge_split_removed"], 0)
        self.assertEqual(counters["inductor"]["scmerge_cat_removed"], 0)
//  [python]
def test_split_cat_merge_mutation(self):
        args = [
            torch.randn(2, 32, 32, 16),
        ]

        def split_cat_mutation(x):
            splits = torch.split(x, 4, dim=1)
            splits[1].copy_(splits[0])
            return torch.cat(splits, dim=1)

        expected = split_cat_mutation(*args)
        actual = torch.compile(split_cat_mutation)(*args)

        torch.testing.assert_close(actual, expected)

        self.assertEqual(counters["inductor"]["scmerge_split_removed"], 0)
        self.assertEqual(counters["inductor"]["scmerge_cat_removed"], 0)
//  [python]
def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
        output = fn()
        baseline_output = baseline_fn()
        accuracy = True
        try:
            t.assert_close(output, baseline_output, atol=1e-2, rtol=0.5)
        except Exception:
            accuracy = False
        finally:
            return accuracy
//  [python]
def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
        output = fn()
        baseline_output = baseline_fn()
        accuracy = True
        try:
            torch.testing.assert_close(output, baseline_output, atol=1e-2, rtol=0.5)
        except Exception:
            accuracy = False
        finally:
            return accuracy
//  [python]
def test_forward(self, inputs):
        model = llama3(
            vocab_size=VOCAB_SIZE,
            num_layers=NUM_LAYERS,
            num_heads=NUM_HEADS,
            num_kv_heads=NUM_KV_HEADS,
            embed_dim=EMBED_DIM,
            max_seq_len=MAX_SEQ_LEN,
        )
        fixed_init_model(model, min_val=-0.25, max_val=0.5)
        actual = model(inputs)
        expected = torch.tensor(3.9763)
        assert actual.shape == (BSZ, SEQ_LEN, VOCAB_SIZE)
        torch.testing.assert_allclose(actual.mean(), expected, atol=1e-4, rtol=1e-4)
//  [python]
def test_forward(self, inputs):
        model = llama3(
            vocab_size=VOCAB_SIZE,
            num_layers=NUM_LAYERS,
            num_heads=NUM_HEADS,
            num_kv_heads=NUM_KV_HEADS,
            embed_dim=EMBED_DIM,
            max_seq_len=MAX_SEQ_LEN,
        )
        fixed_init_model(model, min_val=-0.25, max_val=0.5)
        actual = model(inputs)
        expected = torch.tensor(3.9763)
        assert actual.shape == (BSZ, SEQ_LEN, VOCAB_SIZE)
        torch.testing.assert_close(actual.mean(), expected, atol=1e-4, rtol=1e-4)
//  [python]
def test_recurrent_input_dropout(self):
        drop_prob = 0.5
        encoder = RecurrentEncoder(dropout=drop_prob, emb_dropout=drop_prob)
        input_tensor = torch.Tensor([2, 3, 1, -1])
        encoder.train()
        dropped = encoder.emb_dropout(input=input_tensor)
        # eval switches off dropout
        encoder.eval()
        no_drop = encoder.emb_dropout(input=input_tensor)
        # when dropout is applied, remaining values are divided by drop_prob
        self.assertGreaterEqual((no_drop - (drop_prob * dropped)).abs().sum(), 0)

        drop_prob = 1.0
        encoder = RecurrentEncoder(dropout=drop_prob, emb_dropout=drop_prob)
        all_dropped = encoder.emb_dropout(input=input_tensor)
        self.assertEqual(all_dropped.sum(), 0)
        encoder.eval()
        none_dropped = encoder.emb_dropout(input=input_tensor)
        torch.testing.assert_allclose(no_drop, none_dropped)
        torch.testing.assert_allclose((no_drop - all_dropped), no_drop)
//  [python]
def test_recurrent_input_dropout(self):
        drop_prob = 0.5
        encoder = RecurrentEncoder(dropout=drop_prob, emb_dropout=drop_prob)
        input_tensor = torch.Tensor([2, 3, 1, -1])
        encoder.train()
        dropped = encoder.emb_dropout(input=input_tensor)
        # eval switches off dropout
        encoder.eval()
        no_drop = encoder.emb_dropout(input=input_tensor)
        # when dropout is applied, remaining values are divided by drop_prob
        self.assertGreaterEqual((no_drop - (drop_prob * dropped)).abs().sum(), 0)

        drop_prob = 1.0
        encoder = RecurrentEncoder(dropout=drop_prob, emb_dropout=drop_prob)
        all_dropped = encoder.emb_dropout(input=input_tensor)
        self.assertEqual(all_dropped.sum(), 0)
        encoder.eval()
        none_dropped = encoder.emb_dropout(input=input_tensor)
        torch.testing.assert_close(no_drop, none_dropped)
        torch.testing.assert_close((no_drop - all_dropped), no_drop)
//  [python]
def test_stack_tahn_unbind_merge(self):
        def stack_tahn_unbind(x):
            l1_out = torch.split(x, [20, 20, 20, 10, 10, 20, 20], 1)
            item0 = l1_out[0]
            item1 = l1_out[1]
            item2 = l1_out[2]
            item3 = l1_out[3]
            item4 = l1_out[4]
            item5 = l1_out[5]
            item6 = l1_out[6]
            stack = torch.stack(tensors=(item0, item1, item2), dim=0)
            cat_1 = torch.cat((item3, item4), 1)
            cat_2 = torch.cat((item5, item6), 1)
            tanh = torch.tanh(stack)
            unbind = torch.unbind(tanh, 0)
            return torch.cat((unbind[0], unbind[1], torch.cat((cat_1, cat_2), 1)), 1)

        args = [
            torch.randn(50, 120),
        ]
        for fn, expected_stack_tahn_unbind_merged in [
            (stack_tahn_unbind, 1),
        ]:
            expected = fn(*args)
            actual = torch.compile(fn)(*args)

            self.assertEqual(actual, expected)
            self.assertEqual(
                counters["inductor"]["merge_stack_tahn_unbind_pass"],
                expected_stack_tahn_unbind_merged,
            )
            self.assertIn("merge_getitem_cat_pass_pre_grad", optimus_scuba_log)
            counters.clear()
//  [python]
def test_stack_tahn_unbind_merge(self):
        def stack_tahn_unbind(x):
            l1_out = torch.split(x, [20, 20, 20, 10, 10, 20, 20], 1)
            item0 = l1_out[0]
            item1 = l1_out[1]
            item2 = l1_out[2]
            item3 = l1_out[3]
            item4 = l1_out[4]
            item5 = l1_out[5]
            item6 = l1_out[6]
            stack = torch.stack(tensors=(item0, item1, item2), dim=0)
            cat_1 = torch.cat((item3, item4), 1)
            cat_2 = torch.cat((item5, item6), 1)
            tanh = torch.tanh(stack)
            unbind = torch.unbind(tanh, 0)
            return torch.cat((unbind[0], unbind[1], torch.cat((cat_1, cat_2), 1)), 1)

        args = [
            torch.randn(50, 120),
        ]
        for fn, expected_stack_tahn_unbind_merged in [
            (stack_tahn_unbind, 1),
        ]:
            expected = fn(*args)
            actual = torch.compile(fn)(*args)

            torch.testing.assert_close(actual, expected)
            self.assertEqual(
                counters["inductor"]["merge_stack_tahn_unbind_pass"],
                expected_stack_tahn_unbind_merged,
            )
            self.assertIn("merge_getitem_cat_pass_pre_grad", optimus_scuba_log)
            counters.clear()
//  [python]
def test_against_hf_falcon(kwargs, device, dtype):
    torch.set_default_dtype(dtype)

    ours_config = Config.from_name(**kwargs)
    theirs_config = FalconConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_kv_heads=ours_config.n_query_groups,
        num_hidden_layers=ours_config.n_layer,
        parallel_attn=ours_config.parallel_residual,
        vocab_size=ours_config.padded_vocab_size,
        bias=ours_config.bias,
        new_decoder_architecture=True,
    )

    theirs_model = FalconForCausalLM(theirs_config).to(device)
    theirs_state_dict = theirs_model.state_dict()
    state_dict = {}
    copy_weights_falcon(kwargs["name"], state_dict, theirs_state_dict)
    ours_model = GPT(ours_config).to(device)
    ours_model.load_state_dict(state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32, device=device)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_allclose(ours_y, theirs_y)
//  [python]
def test_against_hf_falcon(kwargs, device, dtype):
    torch.set_default_dtype(dtype)

    ours_config = Config.from_name(**kwargs)
    theirs_config = FalconConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_kv_heads=ours_config.n_query_groups,
        num_hidden_layers=ours_config.n_layer,
        parallel_attn=ours_config.parallel_residual,
        vocab_size=ours_config.padded_vocab_size,
        bias=ours_config.bias,
        new_decoder_architecture=True,
    )

    theirs_model = FalconForCausalLM(theirs_config).to(device)
    theirs_state_dict = theirs_model.state_dict()
    state_dict = {}
    copy_weights_falcon(kwargs["name"], state_dict, theirs_state_dict)
    ours_model = GPT(ours_config).to(device)
    ours_model.load_state_dict(state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32, device=device)
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"]
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_rope_init_under_fsdp():
    """Check that the rope cache is properly initialized"""
    fabric = Fabric(devices=2, strategy="fsdp", accelerator="cuda")
    fabric.launch()

    with fabric.init_module(empty_init=True):
        model = GPT.from_name("pythia-14m", n_layer=1)
    assert model.cos.device.type == "meta"
    assert model.sin.device.type == "meta"

    model = fabric.setup(model)
    assert model.cos.device.type == "cuda"
    assert model.sin.device.type == "cuda"
    cos, sin = model.rope_cache(device=fabric.device)
    t.assert_close(model.cos, cos)
    t.assert_close(model.sin, sin)
//  [python]
def test_rope_init_under_fsdp():
    """Check that the rope cache is properly initialized"""
    fabric = Fabric(devices=2, strategy="fsdp", accelerator="cuda")
    fabric.launch()

    with fabric.init_module(empty_init=True):
        model = GPT.from_name("pythia-14m", n_layer=1)
    assert model.cos.device.type == "meta"
    assert model.sin.device.type == "meta"

    model = fabric.setup(model)
    assert model.cos.device.type == "cuda"
    assert model.sin.device.type == "cuda"
    cos, sin = model.rope_cache(device=fabric.device)
    torch.testing.assert_close(model.cos, cos)
    torch.testing.assert_close(model.sin, sin)
//  [python]
def test_against_original_open_llama_3b(device, dtype):
    torch.set_default_dtype(dtype)

    ours_config = Config.from_name("open_llama_3b", n_layer=2, n_head=8, n_embd=32, intermediate_size=86)
    T = 5
    theirs_config = LlamaConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    theirs_model = LlamaForCausalLM(theirs_config).to(device)
    theirs_state_dict = theirs_model.state_dict()
    state_dict = {}
    copy_weights_hf_llama(ours_config, {}, state_dict, theirs_state_dict)
    ours_model = GPT(ours_config).to(device)
    ours_model.load_state_dict(state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32, device=device)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"].to(dtype)  # HF converts logits to float
    torch.testing.assert_allclose(ours_y, theirs_y)
//  [python]
def test_against_original_open_llama_3b(device, dtype):
    torch.set_default_dtype(dtype)

    ours_config = Config.from_name("open_llama_3b", n_layer=2, n_head=8, n_embd=32, intermediate_size=86)
    T = 5
    theirs_config = LlamaConfig(
        hidden_size=ours_config.n_embd,
        num_attention_heads=ours_config.n_head,
        num_hidden_layers=ours_config.n_layer,
        intermediate_size=ours_config.intermediate_size,
        max_position_embeddings=T,
    )
    assert ours_config.intermediate_size == theirs_config.intermediate_size

    theirs_model = LlamaForCausalLM(theirs_config).to(device)
    theirs_state_dict = theirs_model.state_dict()
    state_dict = {}
    copy_weights_hf_llama(ours_config, {}, state_dict, theirs_state_dict)
    ours_model = GPT(ours_config).to(device)
    ours_model.load_state_dict(state_dict)

    # test end to end
    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32, device=device)
    assert x.size(1) == T
    ours_y = ours_model(x)
    theirs_y = theirs_model(x)["logits"].to(dtype)  # HF converts logits to float
    torch.testing.assert_close(ours_y, theirs_y)
//  [python]
def test_kv_cache(max_seq_length):
    config = Config(block_size=25, padded_vocab_size=5, n_layer=2, n_head=2, n_embd=8)
    model = GPT(config)
    idx = torch.randint(0, model.config.padded_vocab_size, (1, 5))
    max_new_tokens = 20
    model.max_seq_length = max_seq_length
    model.set_kv_cache(1)

    def generate(logits):
        logits = logits[:, -1:]
        probs = torch.nn.functional.softmax(logits, dim=-1)
        return torch.argmax(probs).unsqueeze(0).unsqueeze(0)

    x_no_cache = idx
    x_cache = idx
    input_pos = torch.arange(0, 5)
    for _ in range(max_new_tokens):
        logits_no_cache = model(x_no_cache[:, -max_seq_length:])
        out_no_cache = generate(logits_no_cache)

        logits_cache = model(x_cache, input_pos)
        out_cache = generate(logits_cache)

        torch.testing.assert_allclose(out_no_cache, out_cache, rtol=0, atol=0)

        x_no_cache = torch.cat((x_no_cache, out_no_cache), dim=1)
        x_cache = out_cache
        input_pos = input_pos[-1:] + 1
//  [python]
def test_kv_cache(max_seq_length):
    config = Config(block_size=25, padded_vocab_size=5, n_layer=2, n_head=2, n_embd=8)
    model = GPT(config)
    idx = torch.randint(0, model.config.padded_vocab_size, (1, 5))
    max_new_tokens = 20
    model.max_seq_length = max_seq_length
    model.set_kv_cache(1)

    def generate(logits):
        logits = logits[:, -1:]
        probs = torch.nn.functional.softmax(logits, dim=-1)
        return torch.argmax(probs).unsqueeze(0).unsqueeze(0)

    x_no_cache = idx
    x_cache = idx
    input_pos = torch.arange(0, 5)
    for _ in range(max_new_tokens):
        logits_no_cache = model(x_no_cache[:, -max_seq_length:])
        out_no_cache = generate(logits_no_cache)

        logits_cache = model(x_cache, input_pos)
        out_cache = generate(logits_cache)

        torch.testing.assert_close(out_no_cache, out_cache, rtol=0, atol=0)

        x_no_cache = torch.cat((x_no_cache, out_no_cache), dim=1)
        x_cache = out_cache
        input_pos = input_pos[-1:] + 1
//  [python]
def test_anchor_generator(self):
        images = torch.rand(2, 3, 10, 10)
        features = self.get_features(images)
        model = AnchorGenerator(self.strides, self.anchor_grids)
        model.eval()
        anchors = model(features)

        expected_grids = torch.tensor([[[[[0.0, 0.0], [1.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]]]])
        expected_shifts = torch.tensor([[[[[6.0, 14.0], [6.0, 14.0]], [[6.0, 14.0], [6.0, 14.0]]]]])

        assert len(anchors) == 2
        assert len(anchors[0]) == len(anchors[1]) == 1
        assert tuple(anchors[0][0].shape) == (1, 1, 2, 2, 2)
        assert tuple(anchors[1][0].shape) == (1, 1, 2, 2, 2)

        torch.testing.assert_allclose(anchors[0][0], expected_grids)
        torch.testing.assert_allclose(anchors[1][0], expected_shifts)
//  [python]
def test_anchor_generator(self):
        images = torch.rand(2, 3, 10, 10)
        features = self.get_features(images)
        model = AnchorGenerator(self.strides, self.anchor_grids)
        model.eval()
        anchors = model(features)

        expected_grids = torch.tensor([[[[[0.0, 0.0], [1.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]]]])
        expected_shifts = torch.tensor([[[[[6.0, 14.0], [6.0, 14.0]], [[6.0, 14.0], [6.0, 14.0]]]]])

        assert len(anchors) == 2
        assert len(anchors[0]) == len(anchors[1]) == 1
        assert tuple(anchors[0][0].shape) == (1, 1, 2, 2, 2)
        assert tuple(anchors[1][0].shape) == (1, 1, 2, 2, 2)

        torch.testing.assert_close(anchors[0][0], expected_grids)
        torch.testing.assert_close(anchors[1][0], expected_shifts)
//  [python]
def test_RotaryEmbedding(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for seq_len in range(1024, 4096, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, 0, 0)
            self.Rope = GPTNeoXRotaryEmbedding(dim, max_position_embeddings, base, cuda)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
            position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
            position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
            cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
            result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
            test = self.RopeOP.forward(x)
            torch.testing.assert_allclose(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_RotaryEmbedding(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for seq_len in range(1024, 4096, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, 0, 0)
            self.Rope = GPTNeoXRotaryEmbedding(dim, max_position_embeddings, base, cuda)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
            position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
            position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
            cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
            result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
            test = self.RopeOP.forward(x)
            torch.testing.assert_close(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def assert_close(
    *tensor_pairs: Tuple[Tensor, Tensor],
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    ) -> None:
    """
    Asserts that the two tensors in each pair of tensor_pairs are close.
    See also torch.testing.assert_close.

    Args:
        *tensor_pairs: Pairs of tensors that are asserted to be close.
        rtol: Relative tolerance. If specified, atol must also be specified.
            Otherwise, it is selected according to the tensors' dtypes.
        atol: Absolute tolerance. If specified, rtol must also be specified.
            Otherwise, it is selected according to the tensors' dtypes.
    """
    for pair in tensor_pairs:
        t(*pair, rtol=rtol, atol=atol, equal_nan=True)
//  [python]
def assert_close(
    *tensor_pairs: Tuple[Tensor, Tensor],
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    ) -> None:
    """
    Asserts that the two tensors in each pair of tensor_pairs are close.
    See also torch.testing.assert_close.

    Args:
        *tensor_pairs: Pairs of tensors that are asserted to be close.
        rtol: Relative tolerance. If specified, atol must also be specified.
            Otherwise, it is selected according to the tensors' dtypes.
        atol: Absolute tolerance. If specified, rtol must also be specified.
            Otherwise, it is selected according to the tensors' dtypes.
    """
    for pair in tensor_pairs:
        torch.testing.assert_close(*pair, rtol=rtol, atol=atol, equal_nan=True)
//  [python]
def test_LinearScaling(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for scalar in range(1, 10, 5):
            for seq_len in range(1024, 4096, 1024):
                self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, scalar, 0, 1)
                self.Rope = GPTNeoXLinearScalingRotaryEmbedding(dim, max_position_embeddings, base, cuda, scalar)
                x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
                position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
                cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
                result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
                test = self.RopeOP.forward(x)
                torch.testing.assert_allclose(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_LinearScaling(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for scalar in range(1, 10, 5):
            for seq_len in range(1024, 4096, 1024):
                self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, scalar, 0, 1)
                self.Rope = GPTNeoXLinearScalingRotaryEmbedding(dim, max_position_embeddings, base, cuda, scalar)
                x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
                position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
                cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
                result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
                test = self.RopeOP.forward(x)
                torch.testing.assert_close(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_DynamicNTKScaling(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")
        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for scalar in range(1, 10, 5):
            for seq_len in range(1024, 4096, 1024):
                self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, scalar, 0, 2)
                self.Rope = GPTNeoXDynamicNTKScalingRotaryEmbedding(dim, max_position_embeddings, base, cuda, scalar)
                x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
                position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
                cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
                result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
                test = self.RopeOP.forward(x)
                torch.testing.assert_allclose(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_DynamicNTKScaling(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")
        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for scalar in range(1, 10, 5):
            for seq_len in range(1024, 4096, 1024):
                self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, scalar, 0, 2)
                self.Rope = GPTNeoXDynamicNTKScalingRotaryEmbedding(dim, max_position_embeddings, base, cuda, scalar)
                x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
                position_ids = torch.arange(0, seq_len, dtype=torch.long).cuda()
                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
                cos, sin = self.Rope(x.permute(0, 2, 1, 3), seq_len)
                result, _ = apply_rotary_pos_emb(x.permute(0, 2, 1, 3), x.permute(0, 2, 1, 3), cos, sin, position_ids)
                test = self.RopeOP.forward(x)
                torch.testing.assert_close(test, result.permute(0, 2, 1, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_load_save_without_tied_weights(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.text_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.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.text_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_QWen(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        max_logn_seq_len = 2048
        for seq_len in range(1024, 10000, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, max_logn_seq_len, 3)
            self.Rope = QWenRotaryEmbedding(dim,base)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()

            context_value = math.log(seq_len / max_logn_seq_len, 2) + 1
            ntk_alpha = 2 ** math.ceil(context_value) - 1
            ntk_alpha = max(ntk_alpha, 1)
            freq= self.Rope(seq_len, ntk_alpha = ntk_alpha)
            result = qwen_apply_rotary_pos_emb(x, freq.cuda())
            test = self.RopeOP.forward(x)
            torch.testing.assert_allclose(test, result, atol=1e-2, rtol=1e-2)
//  [python]
def test_QWen(self):
        torch.classes.load_library(os.environ['TEST_SRCDIR'] + "/maga_transformer/tests/libtest_ops.so")

        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        max_logn_seq_len = 2048
        for seq_len in range(1024, 10000, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, max_logn_seq_len, 3)
            self.Rope = QWenRotaryEmbedding(dim,base)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()

            context_value = math.log(seq_len / max_logn_seq_len, 2) + 1
            ntk_alpha = 2 ** math.ceil(context_value) - 1
            ntk_alpha = max(ntk_alpha, 1)
            freq= self.Rope(seq_len, ntk_alpha = ntk_alpha)
            result = qwen_apply_rotary_pos_emb(x, freq.cuda())
            test = self.RopeOP.forward(x)
            torch.testing.assert_close(test, result, atol=1e-2, rtol=1e-2)
//  [python]
def test_glm(self):
        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for seq_len in range(30, 4096, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, 0, 4)
            self.Rope = GLMRotaryEmbedding(dim/2, base)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
            x1, x2 = x.chunk(2, dim=(x.ndim - 1))
            position_ids = torch.arange(0, seq_len, dtype=torch.long).unsqueeze(0).repeat(batch_size, 1).transpose(0,1).cuda()
            cos, sin = self.Rope(x.permute(1, 0, 2, 3), seq_len = seq_len)
            result1, _ = GLM_apply_rotary_pos_emb_index(x1.permute(1, 0, 2, 3), x1.permute(1, 0, 2, 3), cos, sin, position_ids)
            result2, _ = GLM_apply_rotary_pos_emb_index(x2.permute(1, 0, 2, 3), x2.permute(1, 0, 2, 3), cos, sin, position_ids)
            result = torch.concat([result1, result2], dim=(result1.ndim - 1))
            test = self.RopeOP.forward(x)
            torch.testing.assert_allclose(test, result.permute(1, 0, 2, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_glm(self):
        base = 10000
        dim = 128
        headsize = 32
        batch_size = 16
        max_position_embeddings = 2048
        cuda = torch.device('cuda')
        for seq_len in range(30, 4096, 1024):
            self.RopeOP = torch.classes.unittest.RotaryPositionEmbeddingOp(dim, max_position_embeddings, base, 1.0, 0, 4)
            self.Rope = GLMRotaryEmbedding(dim/2, base)
            x = torch.rand(batch_size, seq_len, headsize, dim).cuda()
            x1, x2 = x.chunk(2, dim=(x.ndim - 1))
            position_ids = torch.arange(0, seq_len, dtype=torch.long).unsqueeze(0).repeat(batch_size, 1).transpose(0,1).cuda()
            cos, sin = self.Rope(x.permute(1, 0, 2, 3), seq_len = seq_len)
            result1, _ = GLM_apply_rotary_pos_emb_index(x1.permute(1, 0, 2, 3), x1.permute(1, 0, 2, 3), cos, sin, position_ids)
            result2, _ = GLM_apply_rotary_pos_emb_index(x2.permute(1, 0, 2, 3), x2.permute(1, 0, 2, 3), cos, sin, position_ids)
            result = torch.concat([result1, result2], dim=(result1.ndim - 1))
            test = self.RopeOP.forward(x)
            torch.testing.assert_close(test, result.permute(1, 0, 2, 3), atol=1e-3, rtol=1e-5)
//  [python]
def test_cat_nhwc(self, X, relu):
        # X is NHWC
        X, (scale, zero_point, torch_type) = X

        # Tile out X so # channels is > 64
        X = np.repeat(X, 70 / X.shape[3], 3)
        X = torch.from_numpy(np.ascontiguousarray(X))
        Y = X.clone()
        Y = torch.from_numpy(np.ascontiguousarray(Y))
        # Here, we quantize and get quantized tensors in NHWC for both dims and strides. The
        # permute switches it so that the tensor looks like NCHW but it laid out in memory as
        # NHWC.
        qX = torch.quantize_per_tensor(X, scale, zero_point, torch_type).permute([0, 3, 1, 2])
        qY = torch.quantize_per_tensor(Y, scale, zero_point, torch_type).permute([0, 3, 1, 2])

        ref = torch.cat([qX.dequantize(), qY.dequantize()], dim=1)
        if relu:
            ref[ref < 0] = 0.0
        ref = torch.quantize_per_tensor(ref, scale=scale, zero_point=zero_point, dtype=torch_type)

        if relu:
            out = torch.ops.quantized.cat_relu(
                [qX, qY], dim=1, scale=scale, zero_point=zero_point)
        else:
            out = torch.ops.quantized.cat([qX, qY], dim=1, scale=scale, zero_point=zero_point)

        self.assertEqual(out.dequantize().shape,ref.dequantize().shape)
        self.assertNotEqual(out.stride(), sorted(out.stride()))
//  [python]
def test_cos_anneal_lr_continue(self):
        eta_min = 0.1
        T_max = 5
        scheduler = CosineAnnealingLR(self.opt, T_max=T_max, eta_min=eta_min)
        self.opt.step()
        scheduler.step()
        original_lrs = scheduler._last_lr
        new_scheduler = CosineAnnealingLR(
            self.opt, T_max=T_max, eta_min=eta_min, last_epoch=0
        )
        new_lrs = new_scheduler._last_lr
        self.assertEqual(original_lrs, new_lrs)
//  [python]
def test_cos_anneal_lr_continue(self):
        eta_min = 0.1
        T_max = 5
        scheduler = CosineAnnealingLR(self.opt, T_max=T_max, eta_min=eta_min)
        self.opt.step()
        scheduler.step()
        original_lrs = scheduler._last_lr
        new_scheduler = CosineAnnealingLR(
            self.opt, T_max=T_max, eta_min=eta_min, last_epoch=0
        )
        new_lrs = new_scheduler._last_lr
        torch.testing.assert_close(original_lrs, new_lrs, rtol=1e-4, atol=1e-5)
//  [python]
def rollout_consistency_assertion(
    rollout, *, done_key="done", observation_key="observation", done_strict=False
):
    """Tests that observations in "next" match observations in the next root tensordict when done is False, and don't match otherwise."""

    done = rollout[..., :-1]["next", done_key].squeeze(-1)
    # data resulting from step, when it's not done
    r_not_done = rollout[..., :-1]["next"][~done]
    # data resulting from step, when it's not done, after step_mdp
    r_not_done_tp1 = rollout[:, 1:][~done]
    t.assert_close(
        r_not_done[observation_key],
        r_not_done_tp1[observation_key],
        msg=f"Key {observation_key} did not match",
    )

    if done_strict and not done.any():
        raise RuntimeError("No done detected, test could not complete.")
    if done.any():
        # data resulting from step, when it's done
        r_done = rollout[..., :-1]["next"][done]
        # data resulting from step, when it's done, after step_mdp and reset
        r_done_tp1 = rollout[..., 1:][done]
        # check that at least one obs after reset does not match the version before reset
        assert not torch.isclose(
            r_done[observation_key], r_done_tp1[observation_key]
        ).all()
//  [python]
def rollout_consistency_assertion(
    rollout, *, done_key="done", observation_key="observation", done_strict=False
):
    """Tests that observations in "next" match observations in the next root tensordict when done is False, and don't match otherwise."""

    done = rollout[..., :-1]["next", done_key].squeeze(-1)
    # data resulting from step, when it's not done
    r_not_done = rollout[..., :-1]["next"][~done]
    # data resulting from step, when it's not done, after step_mdp
    r_not_done_tp1 = rollout[:, 1:][~done]
    torch.testing.assert_close(
        r_not_done[observation_key],
        r_not_done_tp1[observation_key],
        msg=f"Key {observation_key} did not match",
    )

    if done_strict and not done.any():
        raise RuntimeError("No done detected, test could not complete.")
    if done.any():
        # data resulting from step, when it's done
        r_done = rollout[..., :-1]["next"][done]
        # data resulting from step, when it's done, after step_mdp and reset
        r_done_tp1 = rollout[..., 1:][done]
        # check that at least one obs after reset does not match the version before reset
        assert not torch.isclose(
            r_done[observation_key], r_done_tp1[observation_key]
        ).all()
//  [python]
def test_backward(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
    _, fk_impl, *_ = ForwardKinematicsFactory(robot, selected_links)
    fk, *_ = get_forward_kinematics_fns(robot, selected_links)

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    jacs_impl = torch.autograd.functional.jacobian(fk_impl, angles, vectorize=True)
    jacs = torch.autograd.functional.jacobian(fk, angles, vectorize=True)

    for jac, jac_impl in zip(jacs, jacs_impl):
        torch.testing.assert_allclose(
            actual=jac_impl, expected=jac, atol=TEST_EPS, rtol=1e-5
        )

    grads = []
    for func in [fk, fk_impl]:
        temp = angles.clone()
        temp.requires_grad = True
        loss = torch.tensor(0, dtype=dtype)
        for pose in func(temp):
            loss = loss + (pose**2).sum()
        loss.backward()
        grads.append(temp.grad)
    torch.testing.assert_allclose(actual=grads[0], expected=grads[1], atol=1e-6, rtol=1e-5)
//  [python]
def test_backward(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
    _, fk_impl, *_ = ForwardKinematicsFactory(robot, selected_links)
    fk, *_ = get_forward_kinematics_fns(robot, selected_links)

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    jacs_impl = torch.autograd.functional.jacobian(fk_impl, angles, vectorize=True)
    jacs = torch.autograd.functional.jacobian(fk, angles, vectorize=True)

    for jac, jac_impl in zip(jacs, jacs_impl):
        torch.testing.assert_close(
            actual=jac_impl, expected=jac, atol=TEST_EPS, rtol=1e-5
        )

    grads = []
    for func in [fk, fk_impl]:
        temp = angles.clone()
        temp.requires_grad = True
        loss = torch.tensor(0, dtype=dtype)
        for pose in func(temp):
            loss = loss + (pose**2).sum()
        loss.backward()
        grads.append(temp.grad)
    torch.testing.assert_close(actual=grads[0], expected=grads[1], atol=1e-6, rtol=1e-5)
//  [python]
def test_embedding_bag_2d_indices(self):
        """
        Tests the case where 2D indices are passed into the operator
        In this case the operator computes the correct offsets argument.
        """
        indices = torch.tensor([[9, 6, 5, 7, 8, 8, 9, 2, 8, 6, 6, 9, 1, 6, 8, 8], [3, 2, 3, 6, 3, 6, 5, 7, 0, 8, 4, 6, 5, 8, 2, 3]])
        weights = torch.randn(10, 12, dtype=torch.float32)

        embedding_bag = torch.nn.EmbeddingBag(
            num_embeddings=10,
            embedding_dim=12,
            include_last_offset=False, _weight=weights,
            scale_grad_by_freq=False, mode='sum'
        )
        result = embedding_bag(indices)

        pt_op = torch.ops.quantized.embedding_bag_byte_rowwise_offsets
        pt_prepack_op = torch.ops.quantized.embedding_bag_byte_prepack
        q_weights = pt_prepack_op(weights)
        qresult = pt_op(q_weights, indices, mode=0, pruned_weights=False)
        self.assertEqual(result, qresult)

        # Test TorchBind based embedding_bag operator
        obs = PerChannelMinMaxObserver(dtype=torch.quint8, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0)
        obs(weights)
        # Get the scale and zero point for the weight tensor
        qparams = obs.calculate_qparams()

        # Quantize the weights to 8bits
        qweight = torch.quantize_per_channel(weights, qparams[0], qparams[1], axis=0, dtype=torch.quint8)

        packed_weight = torch.ops.quantized.embedding_bag_prepack(qweight)
        qresult = torch.ops.quantized.embedding_bag_byte(packed_weight, indices, mode=0)

        self.assertEqual(result, qresult)
//  [python]
def test_tensor_constructor(self, func: Callable[[], torch.Tensor]):
    expected = func()

    with torch_xla2.functions.XLAFunctionMode():
      actual = func()
      self.assertIsInstance(actual, torch_xla2.tensor.XLATensor2)

    torch.testing.assert_allclose(torch_xla2.tensor.j2t(actual._elem), expected)
//  [python]
def test_tensor_constructor(self, func: Callable[[], torch.Tensor]):
    expected = func()

    with torch_xla2.functions.XLAFunctionMode():
      actual = func()
      self.assertIsInstance(actual, torch_xla2.tensor.XLATensor2)

    torch.testing.assert_close(torch_xla2.tensor.j2t(actual._elem), expected)
//  [python]
def test_group_norm(self, input, num_groups):
        num_channels = input.size(1)
        model = _TestModel(num_groups, num_channels)
        model = model.eval().cuda()
        dummy_input = torch.zeros_like(input)
        trt_model = module2trt(model, args=[dummy_input])

        with torch.inference_mode():
            gt = model(input)
            out = trt_model(input)
        torch.testing.assert_allclose(out, gt)
//  [python]
def test_group_norm(self, input, num_groups):
        num_channels = input.size(1)
        model = _TestModel(num_groups, num_channels)
        model = model.eval().cuda()
        dummy_input = torch.zeros_like(input)
        trt_model = module2trt(model, args=[dummy_input])

        with torch.inference_mode():
            gt = model(input)
            out = trt_model(input)
        torch.testing.assert_close(out, gt)
//  [python]
def test_vmap_for_jacobian(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_virtual_ee_link"]
    fk, jfk_b, _ = get_forward_kinematics_fns(robot, selected_links)

    def fun(angles):
        jac_b = jfk_b(angles)[0][0]
        return jac_b.sum(dim=(-1, -2)).unsqueeze(-1)

    def fun_vmap(angles):
        return fun(angles.unsqueeze(0)).squeeze(0)

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    sels = range(batch_size)
    grad = torch.autograd.functional.jacobian(fun, angles, vectorize=True)
    grad_vmap = torch.vmap(torch.func.jacrev(fun_vmap))(angles)

    torch.testing.assert_allclose(
        actual=grad[sels, :, sels], expected=grad_vmap, atol=1e-6, rtol=1e-5
    )
//  [python]
def test_vmap_for_jacobian(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_virtual_ee_link"]
    fk, jfk_b, _ = get_forward_kinematics_fns(robot, selected_links)

    def fun(angles):
        jac_b = jfk_b(angles)[0][0]
        return jac_b.sum(dim=(-1, -2)).unsqueeze(-1)

    def fun_vmap(angles):
        return fun(angles.unsqueeze(0)).squeeze(0)

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    sels = range(batch_size)
    grad = torch.autograd.functional.jacobian(fun, angles, vectorize=True)
    grad_vmap = torch.vmap(torch.func.jacrev(fun_vmap))(angles)

    torch.testing.assert_close(
        actual=grad[sels, :, sels], expected=grad_vmap, atol=1e-6, rtol=1e-5
    )
//  [python]
def test_jacobian(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
    fk, jfk_b, _ = get_forward_kinematics_fns(robot, selected_links)

    def fk_vmap(t):
        return tuple(pose.squeeze(0) for pose in fk(t.unsqueeze(0)))

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    jacs_actual, poses = jfk_b(angles)
    jacs_dense = torch.autograd.functional.jacobian(fk, angles, vectorize=True)
    jacs_vmap = torch.vmap(torch.func.jacrev(fk_vmap))(angles)

    jacs_expected = []
    sels = range(batch_size)
    for pose, jac_dense, jac_vmap in zip(poses, jacs_dense, jacs_vmap):
        jac_sparse = jac_dense[sels, :, :, sels]
        torch.testing.assert_allclose(
            actual=jac_vmap, expected=jac_sparse, atol=1e-6, rtol=1e-5
        )
        jac_sparse_t = jac_sparse.transpose(-1, 1).transpose(-1, -2)
        jac_expected = SE3.left_project(pose, jac_sparse_t).transpose(-1, -2)
        jac_expected[:, 3:] *= 0.5
        jacs_expected.append(jac_expected)

    for jac_actual, jac_expected in zip(jacs_actual, jacs_expected):
        torch.testing.assert_allclose(
            actual=jac_actual, expected=jac_expected, atol=1e-6, rtol=1e-5
        )
//  [python]
def test_jacobian(batch_size: int, dtype: torch.dtype):
    robot = Robot.from_urdf_file(urdf_path, dtype)
    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
    fk, jfk_b, _ = get_forward_kinematics_fns(robot, selected_links)

    def fk_vmap(t):
        return tuple(pose.squeeze(0) for pose in fk(t.unsqueeze(0)))

    rng = torch.Generator()
    rng.manual_seed(0)
    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)

    jacs_actual, poses = jfk_b(angles)
    jacs_dense = torch.autograd.functional.jacobian(fk, angles, vectorize=True)
    jacs_vmap = torch.vmap(torch.func.jacrev(fk_vmap))(angles)

    jacs_expected = []
    sels = range(batch_size)
    for pose, jac_dense, jac_vmap in zip(poses, jacs_dense, jacs_vmap):
        jac_sparse = jac_dense[sels, :, :, sels]
        torch.testing.assert_close(
            actual=jac_vmap, expected=jac_sparse, atol=1e-6, rtol=1e-5
        )
        jac_sparse_t = jac_sparse.transpose(-1, 1).transpose(-1, -2)
        jac_expected = SE3.left_project(pose, jac_sparse_t).transpose(-1, -2)
        jac_expected[:, 3:] *= 0.5
        jacs_expected.append(jac_expected)

    for jac_actual, jac_expected in zip(jacs_actual, jacs_expected):
        torch.testing.assert_close(
            actual=jac_actual, expected=jac_expected, atol=1e-6, rtol=1e-5
        )
//  [python]
def test_conti_tensor(self, batched_tensor, seq_len, conti_tensor):
        conti_out = continuous_tensor(batched_tensor, seq_len)
        self.assertEqual(conti_out, conti_tensor)

        batched_out = batch_tensor(conti_tensor, seq_len)
        self.assertEqual(batched_out, batched_tensor)
//  [python]
def test_conti_tensor(self, batched_tensor, seq_len, conti_tensor):
        conti_out = continuous_tensor(batched_tensor, seq_len)
        torch.testing.assert_close(conti_out, conti_tensor)

        batched_out = batch_tensor(conti_tensor, seq_len)
        torch.testing.assert_close(batched_out, batched_tensor)
//  [python]
def test_where_broadcast(input_shapes):
    torch.manual_seed(42)

    cond_shape = input_shapes[0]
    x_shape = input_shapes[1]
    y_shape = input_shapes[2]

    class Model(torch.nn.Module):
        def forward(self, cond, x, y):
            return torch.where(cond, x, y)

    cond = torch.empty(cond_shape).bernoulli_().to(torch.bool)
    x = torch.randn(x_shape)
    y = torch.randn(y_shape)

    cpu_mod = Model()
    ipu_mod = poptorch.inferenceModel(cpu_mod)

    torch.testing.assert_allclose(actual=ipu_mod(cond, x, y),
                               expected=cpu_mod(cond, x, y))