jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        class Linear(torch.nn.Module):
            def __init__(self, weight, bias=None):
                super(Linear, self).__init__()
                self.weight = weight
                self.bias = bias

            def forward(self, x):
                return F.linear(x, self.weight, self.bias)

        class LinearPrePacked(torch.nn.Module):
            def __init__(self, weight, bias=None):
                super(LinearPrePacked, self).__init__()
                self.packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)

            def forward(self, x):
                return torch.ops.prepacked.linear_clamp_run(x, self.packed_weight_bias)

        data_shape = [batch_size] + list(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand((weight_output_dim))
        else:
            bias = None
        scripted_linear = torch.jit.script(Linear(weight, bias))
        scripted_linear_clamp_prepacked = torch.jit.script(LinearPrePacked(weight, bias))
        input_data = torch.rand(data_shape)
        ref_result = scripted_linear(input_data)
        output_linearprepacked = scripted_linear_clamp_prepacked(input_data)
        torch.testing.assert_allclose(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)

        # Serialize the modules and then deserialize
        input_data = torch.rand(data_shape)
        buffer = io.BytesIO()
        torch.jit.save(scripted_linear, buffer)
        buffer.seek(0)
        deserialized_linear = torch.jit.load(buffer)
        buffer = io.BytesIO()
        torch.jit.save(scripted_linear_clamp_prepacked, buffer)
        buffer.seek(0)
        deserialized_linear_clamp_prepacked = torch.jit.load(buffer)
        ref_result = deserialized_linear(input_data)
        output_linearprepacked = deserialized_linear_clamp_prepacked(input_data)
        torch.testing.assert_allclose(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        class Linear(torch.nn.Module):
            def __init__(self, weight, bias=None):
                super(Linear, self).__init__()
                self.weight = weight
                self.bias = bias

            def forward(self, x):
                return F.linear(x, self.weight, self.bias)

        class LinearPrePacked(torch.nn.Module):
            def __init__(self, weight, bias=None):
                super(LinearPrePacked, self).__init__()
                self.packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)

            def forward(self, x):
                return torch.ops.prepacked.linear_clamp_run(x, self.packed_weight_bias)

        data_shape = [batch_size] + list(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand((weight_output_dim))
        else:
            bias = None
        scripted_linear = torch.jit.script(Linear(weight, bias))
        scripted_linear_clamp_prepacked = torch.jit.script(LinearPrePacked(weight, bias))
        input_data = torch.rand(data_shape)
        ref_result = scripted_linear(input_data)
        output_linearprepacked = scripted_linear_clamp_prepacked(input_data)
        torch.testing.assert_close(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)

        # Serialize the modules and then deserialize
        input_data = torch.rand(data_shape)
        buffer = io.BytesIO()
        torch.jit.save(scripted_linear, buffer)
        buffer.seek(0)
        deserialized_linear = torch.jit.load(buffer)
        buffer = io.BytesIO()
        torch.jit.save(scripted_linear_clamp_prepacked, buffer)
        buffer.seek(0)
        deserialized_linear_clamp_prepacked = torch.jit.load(buffer)
        ref_result = deserialized_linear(input_data)
        output_linearprepacked = deserialized_linear_clamp_prepacked(input_data)
        torch.testing.assert_close(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_two_player():
    world = validation.WinnerLoser.initial(device='cuda')
    agent = validation.ProxyAgent()

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

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

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

    expected = torch.tensor([[+1., -1.]], device=world.device)
    torch.testing.assert_close(m.root().v, expected)
//  [python]
def test_depth():
    world = validation.All.initial(length=3, device='cuda')
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=15, noise_eps=0.)

    expected = torch.tensor([[1/8.]], device=world.device)
    torch.testing.assert_allclose(m.root().v, expected)
//  [python]
def test_depth():
    world = validation.All.initial(length=3, device='cuda')
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=15, noise_eps=0.)

    expected = torch.tensor([[1/8.]], device=world.device)
    torch.testing.assert_close(m.root().v, expected)
//  [python]
def test_multienv():
    # Need to use a fairly complex env here to make sure we've not got 
    # any singleton dims hanging around internally. They can really ruin
    # a tester's day. 
    world = validation.All.initial(n_envs=2, length=3)
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=15, noise_eps=0.)

    expected = torch.tensor([[1/8.], [1/8.]], device=world.device)
    torch.testing.assert_allclose(m.root().v, expected)
//  [python]
def test_multienv():
    # Need to use a fairly complex env here to make sure we've not got 
    # any singleton dims hanging around internally. They can really ruin
    # a tester's day. 
    world = validation.All.initial(n_envs=2, length=3)
    agent = validation.ProxyAgent()

    m = mcts(world, agent, n_nodes=15, noise_eps=0.)

    expected = torch.tensor([[1/8.], [1/8.]], device=world.device)
    torch.testing.assert_close(m.root().v, expected)
//  [python]
def validate_transform_conv1d_to_conv2d(
            self,
            pattern_count_transformed_map,
            pattern_count_optimized_map,
            data_shape):
        input_data = torch.normal(1, 20, size=data_shape)

        for jit_method in ["script", "trace"]:
            module_instance = self
            if jit_method == "script":
                scripted_model = torch.jit.script(module_instance)
            else:
                scripted_model = torch.jit.trace(module_instance, input_data)
            scripted_model.eval()
            ref_result = scripted_model(input_data)
            torch._C._jit_pass_transform_conv1d_to_conv2d(scripted_model._c)
            optimized_scripted_model = optimize_for_mobile(scripted_model)

            buffer = io.BytesIO()
            torch.jit.save(scripted_model, buffer)
            buffer.seek(0)
            deserialized_scripted_model = torch.jit.load(buffer)

            for pattern, v in pattern_count_transformed_map.items():
                if (v == 0):
                    FileCheck().check(pattern).run(deserialized_scripted_model.graph)
                elif (v == -1):
                    FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
                else:
                    FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
            transformed_result = deserialized_scripted_model(input_data)
            torch.testing.assert_allclose(ref_result, transformed_result, rtol=1e-2, atol=1e-3)

            optimized_buffer = io.BytesIO()
            torch.jit.save(optimized_scripted_model, optimized_buffer)
            optimized_buffer.seek(0)
            deserialized_optimized_scripted_model = torch.jit.load(optimized_buffer)

            for pattern, v in pattern_count_optimized_map.items():
                if (v == 0):
                    FileCheck().check(pattern).run(deserialized_optimized_scripted_model.graph)
                elif (v == -1):
                    FileCheck().check_not(pattern).run(deserialized_optimized_scripted_model.graph)
                else:
                    FileCheck().check_count(pattern, v, exactly=True).run(deserialized_optimized_scripted_model.graph)
            xnnpack_result = deserialized_optimized_scripted_model(input_data)
            torch.testing.assert_allclose(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def validate_transform_conv1d_to_conv2d(
            self,
            pattern_count_transformed_map,
            pattern_count_optimized_map,
            data_shape):
        input_data = torch.normal(1, 20, size=data_shape)

        for jit_method in ["script", "trace"]:
            module_instance = self
            if jit_method == "script":
                scripted_model = torch.jit.script(module_instance)
            else:
                scripted_model = torch.jit.trace(module_instance, input_data)
            scripted_model.eval()
            ref_result = scripted_model(input_data)
            torch._C._jit_pass_transform_conv1d_to_conv2d(scripted_model._c)
            optimized_scripted_model = optimize_for_mobile(scripted_model)

            buffer = io.BytesIO()
            torch.jit.save(scripted_model, buffer)
            buffer.seek(0)
            deserialized_scripted_model = torch.jit.load(buffer)

            for pattern, v in pattern_count_transformed_map.items():
                if (v == 0):
                    FileCheck().check(pattern).run(deserialized_scripted_model.graph)
                elif (v == -1):
                    FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
                else:
                    FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
            transformed_result = deserialized_scripted_model(input_data)
            torch.testing.assert_close(ref_result, transformed_result, rtol=1e-2, atol=1e-3)

            optimized_buffer = io.BytesIO()
            torch.jit.save(optimized_scripted_model, optimized_buffer)
            optimized_buffer.seek(0)
            deserialized_optimized_scripted_model = torch.jit.load(optimized_buffer)

            for pattern, v in pattern_count_optimized_map.items():
                if (v == 0):
                    FileCheck().check(pattern).run(deserialized_optimized_scripted_model.graph)
                elif (v == -1):
                    FileCheck().check_not(pattern).run(deserialized_optimized_scripted_model.graph)
                else:
                    FileCheck().check_count(pattern, v, exactly=True).run(deserialized_optimized_scripted_model.graph)
            xnnpack_result = deserialized_optimized_scripted_model(input_data)
            torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def test_leaky_relu(self):
        s = torch.randn(5, 5)
        tg = torch.jit.script(nn.LeakyReLU(0.1))
        o_ref = tg(s)
        tg_a = StaticRuntime(tg)
        o_test = tg_a(s)[0]
        torch.testing.assert_allclose(o_ref, o_test)
//  [python]
def test_leaky_relu(self):
        s = torch.randn(5, 5)
        tg = torch.jit.script(nn.LeakyReLU(0.1))
        o_ref = tg(s)
        tg_a = StaticRuntime(tg)
        o_test = tg_a(s)[0]
        torch.testing.assert_close(o_ref, o_test)
//  [python]
def test_fusion_trivial_graph(self):
        s = torch.full((2, 2), 2)
        tg = torch.jit.script(trivial_graph)
        o_ref = tg(s, s, s)
        torch._C._fuse_to_static_runtime(tg.graph)
        assert "StaticSubgraph" in str(tg.graph)
        o_test = tg(s, s, s)
        torch.testing.assert_allclose(o_ref, o_test)
//  [python]
def test_fusion_trivial_graph(self):
        s = torch.full((2, 2), 2)
        tg = torch.jit.script(trivial_graph)
        o_ref = tg(s, s, s)
        torch._C._fuse_to_static_runtime(tg.graph)
        assert "StaticSubgraph" in str(tg.graph)
        o_test = tg(s, s, s)
        torch.testing.assert_close(o_ref, o_test)
//  [python]
def _assert_graphs_equals(g1: tg.Graph, g2: tg.Graph):
    # Check number of nodes and edges
    assert g1.num_nodes == g2.num_nodes
    assert g1.num_edges == g2.num_edges

    # Check node features
    assert has_node_features(g1) == has_node_features(g2)
    if has_node_features(g1) and has_node_features(g2):
        torch.testing.assert_allclose(g1.node_features, g2.node_features)

    # Check edge indexes
    for sender_id_1, receiver_id_1, sender_id_2, receiver_id_2 in \
            zip(g1.senders, g1.receivers, g2.senders, g2.receivers):
        assert sender_id_1 == sender_id_2
        assert receiver_id_1 == receiver_id_2

    # Check edge features
    assert has_edge_features(g1) == has_edge_features(g2)
    if has_edge_features(g1) and has_edge_features(g2):
        torch.testing.assert_allclose(g1.edge_features, g2.edge_features)

    # Check graph features
    assert has_global_features(g1) == has_global_features(g2)
    if has_global_features(g1) and has_global_features(g2):
        torch.testing.assert_allclose(g1.global_features, g2.global_features)
//  [python]
def _assert_graphs_equals(g1: tg.Graph, g2: tg.Graph):
    # Check number of nodes and edges
    assert g1.num_nodes == g2.num_nodes
    assert g1.num_edges == g2.num_edges

    # Check node features
    assert has_node_features(g1) == has_node_features(g2)
    if has_node_features(g1) and has_node_features(g2):
        torch.testing.assert_close(g1.node_features, g2.node_features)

    # Check edge indexes
    for sender_id_1, receiver_id_1, sender_id_2, receiver_id_2 in \
            zip(g1.senders, g1.receivers, g2.senders, g2.receivers):
        assert sender_id_1 == sender_id_2
        assert receiver_id_1 == receiver_id_2

    # Check edge features
    assert has_edge_features(g1) == has_edge_features(g2)
    if has_edge_features(g1) and has_edge_features(g2):
        torch.testing.assert_close(g1.edge_features, g2.edge_features)

    # Check graph features
    assert has_global_features(g1) == has_global_features(g2)
    if has_global_features(g1) and has_global_features(g2):
        torch.testing.assert_close(g1.global_features, g2.global_features)
//  [python]
def test_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        with torch.no_grad():
            src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

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

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        o_test_kw = attention_a(src, src, value=src, mask=src_mask)

        for a, b in zip(o_ref, o_test):
            torch.testing.assert_allclose(a, b)

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

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

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        o_test_kw = attention_a(src, src, value=src, mask=src_mask)

        for a, b in zip(o_ref, o_test):
            torch.testing.assert_close(a, b)

        for a, b in zip(o_ref, o_test_kw):
            torch.testing.assert_close(a, b)
//  [python]
def _assert_graphs_nx_equal(g1: nx.Graph, g2: nx.Graph):
    # Check number of nodes and edges
    assert g1.number_of_nodes() == g2.number_of_nodes()
    assert g1.number_of_edges() == g2.number_of_edges()

    # Check node features
    for (node_id_1, node_features_1), (node_id_2, node_features_2) in \
            zip(g1.nodes(data='features'), g2.nodes(data='features')):
        assert node_id_1 == node_id_2
        assert (node_features_1 is not None) == (node_features_2 is not None)
        if node_features_1 is not None and node_features_2 is not None:
            torch.testing.assert_allclose(node_features_1, node_features_2)

    # Check edge features
    for (sender_id_1, receiver_id_1, edge_features_1), (sender_id_2, receiver_id_2, edge_features_2) in \
            zip(g1.edges(data='features'), g2.edges(data='features')):
        assert sender_id_1 == sender_id_2
        assert receiver_id_1 == receiver_id_2
        assert (edge_features_1 is not None) == (edge_features_2 is not None)
        if edge_features_1 is not None and edge_features_2 is not None:
            torch.testing.assert_allclose(edge_features_1, edge_features_2)

    # Check graph features
    assert has_global_features(g1) == has_global_features(g2)
    if has_global_features(g1) and has_global_features(g2):
        torch.testing.assert_allclose(g1.graph['features'], g2.graph['features'])
//  [python]
def _assert_graphs_nx_equal(g1: nx.Graph, g2: nx.Graph):
    # Check number of nodes and edges
    assert g1.number_of_nodes() == g2.number_of_nodes()
    assert g1.number_of_edges() == g2.number_of_edges()

    # Check node features
    for (node_id_1, node_features_1), (node_id_2, node_features_2) in \
            zip(g1.nodes(data='features'), g2.nodes(data='features')):
        assert node_id_1 == node_id_2
        assert (node_features_1 is not None) == (node_features_2 is not None)
        if node_features_1 is not None and node_features_2 is not None:
            torch.testing.assert_close(node_features_1, node_features_2)

    # Check edge features
    for (sender_id_1, receiver_id_1, edge_features_1), (sender_id_2, receiver_id_2, edge_features_2) in \
            zip(g1.edges(data='features'), g2.edges(data='features')):
        assert sender_id_1 == sender_id_2
        assert receiver_id_1 == receiver_id_2
        assert (edge_features_1 is not None) == (edge_features_2 is not None)
        if edge_features_1 is not None and edge_features_2 is not None:
            torch.testing.assert_close(edge_features_1, edge_features_2)

    # Check graph features
    assert has_global_features(g1) == has_global_features(g2)
    if has_global_features(g1) and has_global_features(g2):
        torch.testing.assert_close(g1.graph['features'], g2.graph['features'])
//  [python]
def _assert_graph_and_graph_nx_equals(graph: tg.Graph, graph_nx: nx.Graph):
    # Check number of nodes and edges
    assert graph_nx.number_of_nodes() == graph.num_nodes
    assert graph_nx.number_of_edges() == graph.num_edges

    # Check node features
    assert has_node_features(graph) == has_node_features(graph_nx)
    if has_node_features(graph) and has_node_features(graph_nx):
        for node_features_nx, node_features in zip(graph_nx.nodes(data='features'), graph.node_features):
            torch.testing.assert_allclose(node_features_nx[1], node_features)

    # Check edge indexes
    for (sender_id_nx, receiver_id_nx, *_), sender_id, receiver_id in \
            zip(graph_nx.edges, graph.senders, graph.receivers):
        assert sender_id_nx == sender_id
        assert receiver_id_nx == receiver_id

    assert has_edge_features(graph) == has_edge_features(graph_nx)
    if has_edge_features(graph) and has_edge_features(graph_nx):
        for (*_, edge_features_nx), edge_features in zip(graph_nx.edges(data='features'), graph.edge_features):
            torch.testing.assert_allclose(edge_features_nx, edge_features)

    # Check graph features
    assert has_global_features(graph) == has_global_features(graph_nx)
    if has_global_features(graph) and has_global_features(graph_nx):
        torch.testing.assert_allclose(graph_nx.graph['features'], graph.global_features)
//  [python]
def _assert_graph_and_graph_nx_equals(graph: tg.Graph, graph_nx: nx.Graph):
    # Check number of nodes and edges
    assert graph_nx.number_of_nodes() == graph.num_nodes
    assert graph_nx.number_of_edges() == graph.num_edges

    # Check node features
    assert has_node_features(graph) == has_node_features(graph_nx)
    if has_node_features(graph) and has_node_features(graph_nx):
        for node_features_nx, node_features in zip(graph_nx.nodes(data='features'), graph.node_features):
            torch.testing.assert_close(node_features_nx[1], node_features)

    # Check edge indexes
    for (sender_id_nx, receiver_id_nx, *_), sender_id, receiver_id in \
            zip(graph_nx.edges, graph.senders, graph.receivers):
        assert sender_id_nx == sender_id
        assert receiver_id_nx == receiver_id

    assert has_edge_features(graph) == has_edge_features(graph_nx)
    if has_edge_features(graph) and has_edge_features(graph_nx):
        for (*_, edge_features_nx), edge_features in zip(graph_nx.edges(data='features'), graph.edge_features):
            torch.testing.assert_close(edge_features_nx, edge_features)

    # Check graph features
    assert has_global_features(graph) == has_global_features(graph_nx)
    if has_global_features(graph) and has_global_features(graph_nx):
        torch.testing.assert_close(graph_nx.graph['features'], graph.global_features)
//  [python]
def test_fusion_loop(self):
        a = torch.randn(5, 5)
        b = torch.randn(5, 5)
        c = 4
        lg = torch.jit.script(loop_graph)
        o_ref = lg(a, b, c)
        torch._C._fuse_to_static_runtime(lg.graph)
        assert "StaticSubgraph" in str(lg.graph)
        o_test = lg(a, b, c)
        torch.testing.assert_allclose(o_ref, o_test)
//  [python]
def test_fusion_loop(self):
        a = torch.randn(5, 5)
        b = torch.randn(5, 5)
        c = 4
        lg = torch.jit.script(loop_graph)
        o_ref = lg(a, b, c)
        torch._C._fuse_to_static_runtime(lg.graph)
        assert "StaticSubgraph" in str(lg.graph)
        o_test = lg(a, b, c)
        torch.testing.assert_close(o_ref, o_test)
//  [python]
def test_fusion_outputs(self):
        a = torch.randn(2, 2)
        b = torch.randn(2, 2)
        c = 4
        og = torch.jit.script(output_graph)
        o_ref = og(a, b, b, c)
        torch._C._fuse_to_static_runtime(og.graph)
        assert "StaticSubgraph" in str(og.graph)
        o_test = og(a, b, b, c)
        for i in o_ref.keys():
            torch.testing.assert_allclose(o_ref[i], o_test[i])
//  [python]
def test_fusion_outputs(self):
        a = torch.randn(2, 2)
        b = torch.randn(2, 2)
        c = 4
        og = torch.jit.script(output_graph)
        o_ref = og(a, b, b, c)
        torch._C._fuse_to_static_runtime(og.graph)
        assert "StaticSubgraph" in str(og.graph)
        o_test = og(a, b, b, c)
        for i in o_ref.keys():
            torch.testing.assert_close(o_ref[i], o_test[i])
//  [python]
def test_fusion_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        with torch.no_grad():
            src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

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

        torch._C._fuse_to_static_runtime(attention._c)
        o_test = attention(src, src, src, src_mask)

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

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

        torch._C._fuse_to_static_runtime(attention._c)
        o_test = attention(src, src, src, src_mask)

        for a, b in zip(o_ref, o_test):
            torch.testing.assert_close(a, b)
//  [python]
def test_identity_tensor(self):
        h, w = 5, 6
        m = AdaptiveAvgPool2d(output_size=(h, w))
        x = torch.randn(2, 3, h, w, requires_grad=True)
        y = m(x)
        torch.testing.assert_allclose(y, x)
        (dx,) = torch.autograd.grad(y.sum(), inputs=(x,))
        torch.testing.assert_allclose(torch.ones(2, 3, h, w), dx)
//  [python]
def test_identity_tensor(self):
        h, w = 5, 6
        m = AdaptiveAvgPool2d(output_size=(h, w))
        x = torch.randn(2, 3, h, w, requires_grad=True)
        y = m(x)
        torch.testing.assert_close(y, x)
        (dx,) = torch.autograd.grad(y.sum(), inputs=(x,))
        torch.testing.assert_close(torch.ones(2, 3, h, w), dx)
//  [python]
def test_identity_tensor(self):
        h, w = 5, 6
        m = AdaptiveMaxPool2d(output_size=(h, w))
        x = torch.randn(2, 3, h, w, requires_grad=True)
        y = m(x)
        torch.testing.assert_allclose(y, x)
        (dx,) = torch.autograd.grad(y.sum(), inputs=x)
        torch.testing.assert_allclose(dx, torch.ones(2, 3, h, w))
//  [python]
def test_identity_tensor(self):
        h, w = 5, 6
        m = AdaptiveMaxPool2d(output_size=(h, w))
        x = torch.randn(2, 3, h, w, requires_grad=True)
        y = m(x)
        torch.testing.assert_close(y, x)
        (dx,) = torch.autograd.grad(y.sum(), inputs=x)
        torch.testing.assert_close(dx, torch.ones(2, 3, h, w))
//  [python]
def test_forward_tensor(self):
        # N x C x H x W
        m = AdaptiveAvgPool2d(output_size=(1, 2))
        y = m(self.x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[1 / 6, -9 / 6]],
                    # c = 1
                    [[18 / 6, 39 / 6]],
                ],
                # n = 1
                [
                    # c = 0
                    [[24 / 6, 16 / 6]],
                    # c = 1
                    [[20 / 6, 12 / 6]],
                ],
            ],
            dtype=self.x.dtype,
        )
        self.assertEqual(expected_y.size(), y.size())
        torch.testing.assert_allclose(y, expected_y)
        # Check against PyTorch's adaptive pooling
        y2 = torch.nn.functional.adaptive_avg_pool2d(self.x, output_size=(1, 2))
        torch.testing.assert_allclose(y, y2)
//  [python]
def test_forward_tensor(self):
        # N x C x H x W
        m = AdaptiveAvgPool2d(output_size=(1, 2))
        y = m(self.x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[1 / 6, -9 / 6]],
                    # c = 1
                    [[18 / 6, 39 / 6]],
                ],
                # n = 1
                [
                    # c = 0
                    [[24 / 6, 16 / 6]],
                    # c = 1
                    [[20 / 6, 12 / 6]],
                ],
            ],
            dtype=self.x.dtype,
        )
        self.assertEqual(expected_y.size(), y.size())
        torch.testing.assert_close(y, expected_y)
        # Check against PyTorch's adaptive pooling
        y2 = torch.nn.functional.adaptive_avg_pool2d(self.x, output_size=(1, 2))
        torch.testing.assert_close(y, y2)
//  [python]
def test_forward_padded_tensor(self):
        m = AdaptiveAvgPool2d(output_size=(1, 2))
        x = PaddedTensor(self.x, torch.tensor([[2, 2], [1, 3]]))
        y = m(x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[3 / 2, 1 / 2]],
                    # c = 1
                    [[9 / 2, 6 / 2]],
                ],
                # n = 1
                [
                    # c = 0
                    [[12 / 2, 7 / 2]],
                    # c = 1
                    [[2 / 2, 6 / 2]],
                ],
            ],
            dtype=self.x.dtype,
        )
        torch.testing.assert_allclose(y, expected_y)
//  [python]
def test_forward_padded_tensor(self):
        m = AdaptiveAvgPool2d(output_size=(1, 2))
        x = PaddedTensor(self.x, torch.tensor([[2, 2], [1, 3]]))
        y = m(x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[3 / 2, 1 / 2]],
                    # c = 1
                    [[9 / 2, 6 / 2]],
                ],
                # n = 1
                [
                    # c = 0
                    [[12 / 2, 7 / 2]],
                    # c = 1
                    [[2 / 2, 6 / 2]],
                ],
            ],
            dtype=self.x.dtype,
        )
        torch.testing.assert_close(y, expected_y)
//  [python]
def test_eye_cov(self):
        cov_pred = torch.zeros((1, self.covParamNumber))
        cov = self.vec2Cov(cov_pred)
        tt.assert_allclose(cov, torch.eye(3))
//  [python]
def test_eye_cov(self):
        cov_pred = torch.zeros((1, self.covParamNumber))
        cov = self.vec2Cov(cov_pred)
        torch.testing.assert_close(cov, torch.eye(3))
//  [python]
def test_forward_padded_tensor(self):
        m = AdaptiveMaxPool2d(output_size=(1, 2))
        x = PaddedTensor(self.x, torch.tensor([[2, 2], [1, 3]]))
        y = m(x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[2, 1]],
                    # c = 1
                    [[6, 5]],
                ],
                # n = 1
                [
                    # c = 0
                    [[8, 8]],
                    # c = 1
                    [[4, 4]],
                ],
            ],
            dtype=self.x.dtype,
        )
        torch.testing.assert_allclose(y, expected_y)
//  [python]
def test_forward_padded_tensor(self):
        m = AdaptiveMaxPool2d(output_size=(1, 2))
        x = PaddedTensor(self.x, torch.tensor([[2, 2], [1, 3]]))
        y = m(x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[2, 1]],
                    # c = 1
                    [[6, 5]],
                ],
                # n = 1
                [
                    # c = 0
                    [[8, 8]],
                    # c = 1
                    [[4, 4]],
                ],
            ],
            dtype=self.x.dtype,
        )
        torch.testing.assert_close(y, expected_y)
//  [python]
def test_forward_tensor(self):
        m = AdaptiveMaxPool2d(output_size=(1, 2))
        y = m(self.x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[2, 1]],
                    # c = 1
                    [[6, 21]],
                ],
                # n = 1
                [
                    # c = 0
                    [[8, 6]],
                    # c = 1
                    [[9, 6]],
                ],
            ],
            dtype=self.x.dtype,
        )
        self.assertEqual(expected_y.size(), y.size())
        torch.testing.assert_allclose(y, expected_y)
        # Check against PyTorch's adaptive pooling
        y2 = torch.nn.functional.adaptive_max_pool2d(self.x, output_size=(1, 2))
        torch.testing.assert_allclose(y, y2)
//  [python]
def test_forward_tensor(self):
        m = AdaptiveMaxPool2d(output_size=(1, 2))
        y = m(self.x)
        expected_y = torch.tensor(
            [
                # n = 0
                [
                    # c = 0
                    [[2, 1]],
                    # c = 1
                    [[6, 21]],
                ],
                # n = 1
                [
                    # c = 0
                    [[8, 6]],
                    # c = 1
                    [[9, 6]],
                ],
            ],
            dtype=self.x.dtype,
        )
        self.assertEqual(expected_y.size(), y.size())
        torch.testing.assert_close(y, expected_y)
        # Check against PyTorch's adaptive pooling
        y2 = torch.nn.functional.adaptive_max_pool2d(self.x, output_size=(1, 2))
        torch.testing.assert_close(y, y2)
//  [python]
def test_nondiag_coef(self):
        cov_pred = torch.ones((1, self.covParamNumber))
        cov_pred[0, 3:6] = 1e9
        cov = self.vec2Cov(cov_pred)
        tt.assert_allclose(cov, np.exp(2), atol=1e-3, rtol=1e-3)
//  [python]
def test_nondiag_coef(self):
        cov_pred = torch.ones((1, self.covParamNumber))
        cov_pred[0, 3:6] = 1e9
        cov = self.vec2Cov(cov_pred)
        torch.testing.assert_close(cov, np.exp(2), atol=1e-3, rtol=1e-3)
//  [python]
def test_1_cov_no_error(self):
        target = torch.FloatTensor([[1, 1, 1]])
        mean = torch.FloatTensor([[0, 0, 0]])
        cov_pred = torch.zeros((1, self.covParamNumber))
        cov_pred[0, 0:3] = 1

        cov = self.vec2Cov(cov_pred)
        tt.assert_allclose(cov, np.exp(2) * torch.eye(3))

        expected_err = 3.0 / np.exp(2) + 3 * 2
        err = self.toMahalanobisDistance(target, mean, cov_pred)
        nt.assert_equal(err, torch.FloatTensor([[expected_err]]))
//  [python]
def test_1_cov_no_error(self):
        target = torch.FloatTensor([[1, 1, 1]])
        mean = torch.FloatTensor([[0, 0, 0]])
        cov_pred = torch.zeros((1, self.covParamNumber))
        cov_pred[0, 0:3] = 1

        cov = self.vec2Cov(cov_pred)
        torch.testing.assert_close(cov, np.exp(2) * torch.eye(3))

        expected_err = 3.0 / np.exp(2) + 3 * 2
        err = self.toMahalanobisDistance(target, mean, cov_pred)
        nt.assert_equal(err, torch.FloatTensor([[expected_err]]))
//  [python]
def test_superpos_empty(self):
        stack = TransitionParserStack.empty(1, 4, None)
        policy = torch.ones(1, 3) / 3
        stack.update(SHIFT, VEC1)
        stack.update(SHIFT, VEC2)
        stack.update(policy, VEC1)
        expected = [[
            [2/3, 1/3, 0, 0],
            [0, 1/3, 0, 0],
            [1/3, 0, 0, 0],
        ]]
        torch.testing.assert_allclose(stack.tapes.tolist(), expected)
//  [python]
def test_superpos_empty(self):
        stack = TransitionParserStack.empty(1, 4, None)
        policy = torch.ones(1, 3) / 3
        stack.update(SHIFT, VEC1)
        stack.update(SHIFT, VEC2)
        stack.update(policy, VEC1)
        expected = [[
            [2/3, 1/3, 0, 0],
            [0, 1/3, 0, 0],
            [1/3, 0, 0, 0],
        ]]
        torch.testing.assert_close(stack.tapes.tolist(), expected)
//  [python]
def test_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = self.prepare_inputs()

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cpu() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=40 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0], torch.tensor(self.EXPECTED_OUTPUT_2))

        set_seed(0)
        zs = model._sample(zs, labels, [1], sample_length=40 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[1][0], torch.tensor(self.EXPECTED_OUTPUT_1))

        set_seed(0)
        zs = model._sample(zs, labels, [2], sample_length=40 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[2][0], torch.tensor(self.EXPECTED_OUTPUT_0))
//  [python]
def test_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = self.prepare_inputs()

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cpu() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=40 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0], torch.tensor(self.EXPECTED_OUTPUT_2))

        set_seed(0)
        zs = model._sample(zs, labels, [1], sample_length=40 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[1][0], torch.tensor(self.EXPECTED_OUTPUT_1))

        set_seed(0)
        zs = model._sample(zs, labels, [2], sample_length=40 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[2][0], torch.tensor(self.EXPECTED_OUTPUT_0))
//  [python]
def test_fp16_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval().half().to("cuda")
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
//  [python]
def test_fp16_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval().half().to("cuda")
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
//  [python]
def test_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = self.prepare_inputs(self.model_id)

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cpu() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0], torch.tensor(self.EXPECTED_OUTPUT_2))

        set_seed(0)
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[1][0], torch.tensor(self.EXPECTED_OUTPUT_1))

        set_seed(0)
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[2][0], torch.tensor(self.EXPECTED_OUTPUT_0))
//  [python]
def test_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = self.prepare_inputs(self.model_id)

        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long).cpu() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0], torch.tensor(self.EXPECTED_OUTPUT_2))

        set_seed(0)
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[1][0], torch.tensor(self.EXPECTED_OUTPUT_1))

        set_seed(0)
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[2][0], torch.tensor(self.EXPECTED_OUTPUT_0))
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval().to("cuda")
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].cuda()
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].cuda()
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].cuda()
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval().to("cuda")
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].cuda()
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].cuda()
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].cuda()
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_fc_quant_nobias(quant_mod, batch, seq, in_features, out_features):
    # This test may fail due to the possibility of a difference between the quantize_per_token kernel in the CUDA
    # and the PyTorch, However, this difference is likely to be small and will not significantly affect the result.
    input_tensor = torch.randn((batch, seq, in_features)).half().cuda()
    fc = torch.nn.Linear(in_features=in_features, out_features=out_features, bias=None).half()

    fc_copy = copy.deepcopy(fc)
    quantized_fc_real = tinynn_m.QLinear(fc, quant_mode=quant_mod).cuda()
    q_output = quantized_fc_real(input_tensor)

    tinynn_m.SPEEDUP = False
    quantized_fc_simu = tinynn_m.QLinear(fc_copy, quant_mode=quant_mod).cuda()
    fq_output = quantized_fc_simu(input_tensor)

    torch.testing.assert_allclose(q_output, fq_output)
//  [python]
def test_fc_quant_nobias(quant_mod, batch, seq, in_features, out_features):
    # This test may fail due to the possibility of a difference between the quantize_per_token kernel in the CUDA
    # and the PyTorch, However, this difference is likely to be small and will not significantly affect the result.
    input_tensor = torch.randn((batch, seq, in_features)).half().cuda()
    fc = torch.nn.Linear(in_features=in_features, out_features=out_features, bias=None).half()

    fc_copy = copy.deepcopy(fc)
    quantized_fc_real = tinynn_m.QLinear(fc, quant_mode=quant_mod).cuda()
    q_output = quantized_fc_real(input_tensor)

    tinynn_m.SPEEDUP = False
    quantized_fc_simu = tinynn_m.QLinear(fc_copy, quant_mode=quant_mod).cuda()
    fq_output = quantized_fc_simu(input_tensor)

    torch.testing.assert_close(q_output, fq_output)
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = [i for i in self.prepare_inputs()]

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = [i for i in self.prepare_inputs()]

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_close(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_close(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_close(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_bbox_transform(
        self,
        roi_counts,
        num_classes,
        rotated,
        angle_bound_on,
        clip_angle_thresh,
        gc,
        dc,
    ):
        """
        Test with rois for multiple images in a batch
        """
        rois, deltas, im_info = create_bbox_transform_inputs(
            roi_counts, num_classes, rotated
        )

        def bbox_transform_ref():
            ref_op = core.CreateOperator(
                "BBoxTransform",
                ["rois", "deltas", "im_info"],
                ["box_out"],
                apply_scale=False,
                rotated=rotated,
                angle_bound_on=angle_bound_on,
                clip_angle_thresh=clip_angle_thresh,
            )
            workspace.FeedBlob("rois", rois)
            workspace.FeedBlob("deltas", deltas)
            workspace.FeedBlob("im_info", im_info)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("box_out")

        box_out = torch.tensor(bbox_transform_ref())
        a, b = torch.ops._caffe2.BBoxTransform(
            torch.tensor(rois),
            torch.tensor(deltas),
            torch.tensor(im_info),
            [1.0, 1.0, 1.0, 1.0],
            False,
            rotated,
            angle_bound_on,
            -90,
            90,
            clip_angle_thresh,
        )

        torch.testing.assert_allclose(box_out, a)