jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _test_perm102_bmm_rcr(
        self,
        B=25,
        M=128,
        K=256,
        N=100,
        dtype="float16",
        test_name="perm102_bmm_rcr",
    ):
        target = detect_target()
        X = Tensor(shape=[M, B, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[B, N, K], dtype=dtype, name="input_1", is_input=True)
        OP = ops.perm102_bmm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(Y, target, "./tmp", test_name)

        X_pt = get_random_torch_tensor(shape=(M, B, K), dtype=dtype)
        W_pt = get_random_torch_tensor(shape=(B, N, K), dtype=dtype)

        XT = X_pt.permute(1, 0, 2)
        Y_pt = torch.bmm(XT, W_pt.permute([0, 2, 1]))
        Y_pt = Y_pt.permute(1, 0, 2).contiguous()
        y = torch.empty_like(Y_pt)
        module.run_with_tensors({"input_0": X_pt, "input_1": W_pt}, [y])

        torch.allclose(Y_pt, y, atol=1e-1, rtol=1e-1)
//  [python]
def _test_perm102_bmm_rcr(
        self,
        B=25,
        M=128,
        K=256,
        N=100,
        dtype="float16",
        test_name="perm102_bmm_rcr",
    ):
        target = detect_target()
        X = Tensor(shape=[M, B, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[B, N, K], dtype=dtype, name="input_1", is_input=True)
        OP = ops.perm102_bmm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(Y, target, "./tmp", test_name)

        X_pt = get_random_torch_tensor(shape=(M, B, K), dtype=dtype)
        W_pt = get_random_torch_tensor(shape=(B, N, K), dtype=dtype)

        XT = X_pt.permute(1, 0, 2)
        Y_pt = torch.bmm(XT, W_pt.permute([0, 2, 1]))
        Y_pt = Y_pt.permute(1, 0, 2).contiguous()
        y = torch.empty_like(Y_pt)
        module.run_with_tensors({"input_0": X_pt, "input_1": W_pt}, [y])

        torch.testing.assert_close(Y_pt, y, atol=1e-1, rtol=1e-1)
//  [python]
def test_getAction_epsilon_greedy_distributed(
    epsilon, state_dims, action_dims, discrete_actions
):
    accelerator = Accelerator()
    agent_ids = ["agent_0", "agent_1"]
    state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    from agilerl.algorithms.maddpg import MADDPG

    maddpg = MADDPG(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        accelerator=accelerator,
    )
    new_actors = [
        DummyEvolvableMLP(
            num_inputs=actor.num_inputs,
            num_outputs=actor.num_outputs,
            hidden_size=actor.hidden_size,
            device=actor.device,
            mlp_output_activation=actor.mlp_output_activation,
        )
        for actor in maddpg.actors
    ]
    maddpg.actors = new_actors
    cont_actions, discrete_action = maddpg.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            t.assert_allclose(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
def test_getAction_epsilon_greedy_distributed(
    epsilon, state_dims, action_dims, discrete_actions
):
    accelerator = Accelerator()
    agent_ids = ["agent_0", "agent_1"]
    state = {agent: np.random.randn(*state_dims[0]) for agent in agent_ids}
    from agilerl.algorithms.maddpg import MADDPG

    maddpg = MADDPG(
        state_dims,
        action_dims,
        one_hot=False,
        n_agents=2,
        agent_ids=agent_ids,
        max_action=[[1], [1]],
        min_action=[[-1], [-1]],
        discrete_actions=discrete_actions,
        accelerator=accelerator,
    )
    new_actors = [
        DummyEvolvableMLP(
            num_inputs=actor.num_inputs,
            num_outputs=actor.num_outputs,
            hidden_size=actor.hidden_size,
            device=actor.device,
            mlp_output_activation=actor.mlp_output_activation,
        )
        for actor in maddpg.actors
    ]
    maddpg.actors = new_actors
    cont_actions, discrete_action = maddpg.getAction(state, epsilon)
    if discrete_actions:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            torch.testing.assert_close(sum(action), 1.00)
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
        for idx, action in enumerate(list(discrete_action.values())):
            assert action <= action_dims[idx] - 1
    else:
        for idx, action in enumerate(list(cont_actions.values())):
            assert len(action) == action_dims[idx]
            act = action[idx]
            assert isinstance(act, np.float32)
            assert -1 <= act <= 1
//  [python]
if MAIN:
	attn_patterns_layer_0_copy = gpt2_cache["blocks.0.attn.hook_pattern"]
	
	t.assert_close(attn_patterns_layer_0, attn_patterns_layer_0_copy)
//  [python]
if MAIN:
	print(gpt2_small.to_str_tokens("gpt2"))
	print(gpt2_small.to_tokens("gpt2"))
	print(gpt2_small.to_string([50256, 70, 457, 17]))

# %%


if MAIN:
	logits: Tensor = gpt2_small(model_description_text, return_type="logits")
	prediction = logits.argmax(dim=-1).squeeze()[:-1]
	# FLAT SOLUTION
	# YOUR CODE HERE - get the model's prediction on the text
	true_tokens = gpt2_small.to_tokens(model_description_text).squeeze()[1:]
	is_correct = (prediction == true_tokens)
	
	print(f"Model accuracy: {is_correct.sum()}/{len(true_tokens)}")
	print(f"Correct words: {gpt2_small.to_str_tokens(prediction[is_correct])}")

# %%


if MAIN:
	gpt2_text = "Natural language processing tasks, such as question answering, machine translation, reading comprehension, and summarization, are typically approached with supervised learning on taskspecific datasets."
	gpt2_tokens = gpt2_small.to_tokens(gpt2_text)
	gpt2_logits, gpt2_cache = gpt2_small.run_with_cache(gpt2_tokens, remove_batch_dim=True)

# %%


if MAIN:
	attn_patterns_layer_0 = gpt2_cache["pattern", 0]

# %%


if MAIN:
	attn_patterns_layer_0_copy = gpt2_cache["blocks.0.attn.hook_pattern"]
	
	t.testing.assert_close(attn_patterns_layer_0, attn_patterns_layer_0_copy)
//  [python]
def test_fuser_util(self, partition):
        m = TestModule()
        gm = symbolic_trace(m)

        nodes_by_name = {node.name : node for node in gm.graph.nodes}

        partitions = []
        for node_names in partition:
            partitions.append([nodes_by_name[name] for name in node_names])

        fused_graph = fuse_by_partitions(gm, partitions)

        a, b, c = torch.rand(4), torch.rand(4), torch.rand(4)

        expected = m(a, b, c)
        result = fused_graph(a, b, c)

        self.assertEqual(expected, result)
//  [python]
def test_fuser_util(self, partition):
        m = TestModule()
        gm = symbolic_trace(m)

        nodes_by_name = {node.name : node for node in gm.graph.nodes}

        partitions = []
        for node_names in partition:
            partitions.append([nodes_by_name[name] for name in node_names])

        fused_graph = fuse_by_partitions(gm, partitions)

        a, b, c = torch.rand(4), torch.rand(4), torch.rand(4)

        expected = m(a, b, c)
        result = fused_graph(a, b, c)

        torch.testing.assert_close(expected, result)
//  [python]
def test_partitioner(self, fn, expected_partition, bookend_non_compute_pass):
        traced = symbolic_trace(fn)

        non_compute_ops = []
        if bookend_non_compute_pass:
            non_compute_ops = ["torch.ops.aten.view", "torch.ops.aten.permute"]

        supported_ops = MockOperatorSupport()
        partitioner = CapabilityBasedPartitioner(traced,
                                                 supported_ops,
                                                 allows_single_node_partition=True,
                                                 non_compute_ops=non_compute_ops)
        partitions = partitioner.propose_partitions()
        if bookend_non_compute_pass:
            partitioner.remove_bookend_non_compute_ops(partitions)

        partitions_name = [[node.name for node in partition.nodes] for partition in partitions]
        assert len(partitions_name) == len(expected_partition)
        for i in range(len(partitions_name)):
            assert set(partitions_name[i]) == set(expected_partition[i])

        fused_graph = partitioner.fuse_partitions(partitions)

        a, b, c = torch.rand(4), torch.rand(4), torch.rand(4)

        expected = fn(a, b, c)
        result = fused_graph(a, b, c)
        torch.testing.assert_allclose(expected, result)
//  [python]
def test_partitioner(self, fn, expected_partition, bookend_non_compute_pass):
        traced = symbolic_trace(fn)

        non_compute_ops = []
        if bookend_non_compute_pass:
            non_compute_ops = ["torch.ops.aten.view", "torch.ops.aten.permute"]

        supported_ops = MockOperatorSupport()
        partitioner = CapabilityBasedPartitioner(traced,
                                                 supported_ops,
                                                 allows_single_node_partition=True,
                                                 non_compute_ops=non_compute_ops)
        partitions = partitioner.propose_partitions()
        if bookend_non_compute_pass:
            partitioner.remove_bookend_non_compute_ops(partitions)

        partitions_name = [[node.name for node in partition.nodes] for partition in partitions]
        assert len(partitions_name) == len(expected_partition)
        for i in range(len(partitions_name)):
            assert set(partitions_name[i]) == set(expected_partition[i])

        fused_graph = partitioner.fuse_partitions(partitions)

        a, b, c = torch.rand(4), torch.rand(4), torch.rand(4)

        expected = fn(a, b, c)
        result = fused_graph(a, b, c)
        torch.testing.assert_close(expected, result)
//  [python]
def test_partitioner_independent_output(self, fn, expected_partition):
        traced = symbolic_trace(fn)

        supported_ops = MockOperatorSupport()
        partitioner = CapabilityBasedPartitioner(traced,
                                                 supported_ops,
                                                 allows_single_node_partition=True)
        partitions = partitioner.propose_partitions()
        partitions_name = [[node.name for node in partition.nodes] for partition in partitions]
        assert len(partitions_name) == len(expected_partition)
        for i in range(len(partitions_name)):
            assert set(partitions_name[i]) == set(expected_partition[i])

        fused_graph = partitioner.fuse_partitions(partitions)

        a, b, c, d, e, f = torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4)

        expected = fn(a, b, c, d, e, f)
        result = fused_graph(a, b, c, d, e, f)
        t.assert_close(expected, result)
//  [python]
def test_partitioner_independent_output(self, fn, expected_partition):
        traced = symbolic_trace(fn)

        supported_ops = MockOperatorSupport()
        partitioner = CapabilityBasedPartitioner(traced,
                                                 supported_ops,
                                                 allows_single_node_partition=True)
        partitions = partitioner.propose_partitions()
        partitions_name = [[node.name for node in partition.nodes] for partition in partitions]
        assert len(partitions_name) == len(expected_partition)
        for i in range(len(partitions_name)):
            assert set(partitions_name[i]) == set(expected_partition[i])

        fused_graph = partitioner.fuse_partitions(partitions)

        a, b, c, d, e, f = torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4), torch.rand(4)

        expected = fn(a, b, c, d, e, f)
        result = fused_graph(a, b, c, d, e, f)
        torch.testing.assert_close(expected, result)
//  [python]
def test_expand_into_jagged_permute(
        self,
        T: int,
        W: int,
    ) -> None:
        length_per_w = [random.randint(5000, 10000) for i in range(W)]
        length_1d = list(
            itertools.chain.from_iterable(itertools.repeat(x, T) for x in length_per_w)
        )
        permute_list = list(range(T * W))
        random.shuffle(permute_list)
        permuted_length_1d = [length_1d[r] for r in permute_list]
        permute_tensor = torch.tensor(permute_list)

        # compute offsets
        offsets_1d = [0] + list(itertools.accumulate(length_1d))
        permuted_offsets_1d = [0] + list(itertools.accumulate(permuted_length_1d))
        offsets_1d_tensor = torch.tensor(offsets_1d)
        permuted_offsets_1d_tensor = torch.tensor(permuted_offsets_1d)

        # cpu op
        output_permute_cpu = torch.ops.fbgemm.expand_into_jagged_permute(
            permute_tensor,
            offsets_1d_tensor,
            permuted_offsets_1d_tensor,
            offsets_1d[-1],
        )

        # reference solution
        output_permute_ref = self.expand_into_jagged_permute_ref_(
            permute_list,
            length_1d,
        )
        output_permute_ref_tensor = torch.tensor(output_permute_ref)

        # assert cpu and gpu ops
        torch.testing.assert_allclose(output_permute_cpu, output_permute_ref_tensor)
        if gpu_available:
            # gpu op
            output_permute_gpu = torch.ops.fbgemm.expand_into_jagged_permute(
                permute_tensor.cuda(),
                offsets_1d_tensor.cuda(),
                permuted_offsets_1d_tensor.cuda(),
                offsets_1d[-1],
            )
            torch.testing.assert_allclose(
                output_permute_gpu.cpu(), output_permute_ref_tensor
            )
//  [python]
def test_expand_into_jagged_permute(
        self,
        T: int,
        W: int,
    ) -> None:
        length_per_w = [random.randint(5000, 10000) for i in range(W)]
        length_1d = list(
            itertools.chain.from_iterable(itertools.repeat(x, T) for x in length_per_w)
        )
        permute_list = list(range(T * W))
        random.shuffle(permute_list)
        permuted_length_1d = [length_1d[r] for r in permute_list]
        permute_tensor = torch.tensor(permute_list)

        # compute offsets
        offsets_1d = [0] + list(itertools.accumulate(length_1d))
        permuted_offsets_1d = [0] + list(itertools.accumulate(permuted_length_1d))
        offsets_1d_tensor = torch.tensor(offsets_1d)
        permuted_offsets_1d_tensor = torch.tensor(permuted_offsets_1d)

        # cpu op
        output_permute_cpu = torch.ops.fbgemm.expand_into_jagged_permute(
            permute_tensor,
            offsets_1d_tensor,
            permuted_offsets_1d_tensor,
            offsets_1d[-1],
        )

        # reference solution
        output_permute_ref = self.expand_into_jagged_permute_ref_(
            permute_list,
            length_1d,
        )
        output_permute_ref_tensor = torch.tensor(output_permute_ref)

        # assert cpu and gpu ops
        torch.testing.assert_close(output_permute_cpu, output_permute_ref_tensor)
        if gpu_available:
            # gpu op
            output_permute_gpu = torch.ops.fbgemm.expand_into_jagged_permute(
                permute_tensor.cuda(),
                offsets_1d_tensor.cuda(),
                permuted_offsets_1d_tensor.cuda(),
                offsets_1d[-1],
            )
            torch.testing.assert_close(
                output_permute_gpu.cpu(), output_permute_ref_tensor
            )
//  [python]
def compare_feed_forward(embed_dim: int, hidden_dim: int) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(0)

    # generate input tensor used by both implementations
    input_t = torch.randn(1, embed_dim)

    # reference implementation; initialize with constant to compare outputs
    ff_ref = FeedForwardRef(dim=embed_dim, hidden_dim=4 * embed_dim)
    fixed_init_model(ff_ref)

    with torch.no_grad():
        ff_out_ref = ff_ref(input_t)

    hidden_dim = _scale_hidden_dim_for_mlp(embed_dim)
    ff = FeedForward(dim=embed_dim, hidden_dim=hidden_dim, linear_class=torch.nn.Linear)
    fixed_init_model(ff)

    with torch.no_grad():
        ff_out = ff(input_t)

    torch.testing.assert_allclose(ff_out, ff_out_ref, atol=1e-5, rtol=1e-5)
    print(ff_out.mean())
    print(ff_out.max())
//  [python]
def compare_feed_forward(embed_dim: int, hidden_dim: int) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(0)

    # generate input tensor used by both implementations
    input_t = torch.randn(1, embed_dim)

    # reference implementation; initialize with constant to compare outputs
    ff_ref = FeedForwardRef(dim=embed_dim, hidden_dim=4 * embed_dim)
    fixed_init_model(ff_ref)

    with torch.no_grad():
        ff_out_ref = ff_ref(input_t)

    hidden_dim = _scale_hidden_dim_for_mlp(embed_dim)
    ff = FeedForward(dim=embed_dim, hidden_dim=hidden_dim, linear_class=torch.nn.Linear)
    fixed_init_model(ff)

    with torch.no_grad():
        ff_out = ff(input_t)

    torch.testing.assert_close(ff_out, ff_out_ref, atol=1e-5, rtol=1e-5)
    print(ff_out.mean())
    print(ff_out.max())
//  [python]
def test_spmm_coo_dequant(dim1, dim2, dtype):
    threshold = 6.0
    # threshold = 2.8
    # threshold = 0.0
    A = torch.randn(dim1, dim2, device="cuda").half()
    B = torch.empty(dim2, dim2 * 4, device="cuda", dtype=torch.float16)
    torch.nn.init.xavier_uniform_(B)
    Bt = B.t().contiguous()

    CB, CBt, statsB, statsBt, coo_tensor = F.double_quant(B)

    rowidx = torch.randint(0, A.shape[-1], size=(15,))

    A[:, rowidx] = 8.0

    idx = torch.abs(A) >= threshold
    nnz = (idx == 1).sum().item()
    rows, cols = torch.where(idx)
    values = A[idx]
    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
    A2 = A * idx
    out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
    out1 = torch.matmul(A2, B.half())
    out3 = F.spmm_coo_very_sparse(cooA, CBt.half())
    out3 = out3 * statsBt.half() / 127

    values, counts = torch.unique(cooA.rowidx, return_counts=True)
    offset = counts.cumsum(0).int()
    max_count, max_idx = torch.sort(counts, descending=True)
    print(torch.median(max_count.float()))

    torch.testing.assert_allclose(out2, out3, rtol=0.05, atol=0.001)

    p = 200 / (2048 * 12288 * 4)
    n = out1.numel()
    count = math.ceil(p * n)
    assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=count)

    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(100):
    #   out2 = F.spmm_coo_very_sparse(cooA, B)
    # torch.cuda.synchronize()
    # print('fp16', time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo(cooA, B)
    torch.cuda.synchronize()
    print("cusparse fp16", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo_very_sparse(cooA, CBt)
    torch.cuda.synchronize()
    print("int8", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
    torch.cuda.synchronize()
    print("int8+dequant", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = torch.matmul(A, B)
    torch.cuda.synchronize()
    print("matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
        out = out1 + out2
    torch.cuda.synchronize()
    print("sparse+ matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
        torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1)
    torch.cuda.synchronize()
    print("partial matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
    torch.cuda.synchronize()
    print("partial matmul", time.time() - t0)
//  [python]
def test_spmm_coo_dequant(dim1, dim2, dtype):
    threshold = 6.0
    # threshold = 2.8
    # threshold = 0.0
    A = torch.randn(dim1, dim2, device="cuda").half()
    B = torch.empty(dim2, dim2 * 4, device="cuda", dtype=torch.float16)
    torch.nn.init.xavier_uniform_(B)
    Bt = B.t().contiguous()

    CB, CBt, statsB, statsBt, coo_tensor = F.double_quant(B)

    rowidx = torch.randint(0, A.shape[-1], size=(15,))

    A[:, rowidx] = 8.0

    idx = torch.abs(A) >= threshold
    nnz = (idx == 1).sum().item()
    rows, cols = torch.where(idx)
    values = A[idx]
    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
    A2 = A * idx
    out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
    out1 = torch.matmul(A2, B.half())
    out3 = F.spmm_coo_very_sparse(cooA, CBt.half())
    out3 = out3 * statsBt.half() / 127

    values, counts = torch.unique(cooA.rowidx, return_counts=True)
    offset = counts.cumsum(0).int()
    max_count, max_idx = torch.sort(counts, descending=True)
    print(torch.median(max_count.float()))

    torch.testing.assert_close(out2, out3, rtol=0.05, atol=0.001)

    p = 200 / (2048 * 12288 * 4)
    n = out1.numel()
    count = math.ceil(p * n)
    assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=count)

    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(100):
    #   out2 = F.spmm_coo_very_sparse(cooA, B)
    # torch.cuda.synchronize()
    # print('fp16', time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo(cooA, B)
    torch.cuda.synchronize()
    print("cusparse fp16", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo_very_sparse(cooA, CBt)
    torch.cuda.synchronize()
    print("int8", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
    torch.cuda.synchronize()
    print("int8+dequant", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out2 = torch.matmul(A, B)
    torch.cuda.synchronize()
    print("matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
        out = out1 + out2
    torch.cuda.synchronize()
    print("sparse+ matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
        torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1)
    torch.cuda.synchronize()
    print("partial matmul", time.time() - t0)

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(100):
        out1 = bnb.matmul(A, Bt)
    torch.cuda.synchronize()
    print("partial matmul", time.time() - t0)
//  [python]
def _test_jagged_lengths_to_offsets(
        self,
        batch_size: int,
        offsets_dtype: str = "int32",
        test_suffix: str = "",
    ):
        LENGTHS = Tensor(
            shape=[IntVar([1, batch_size], name="batch_size")],
            name="lengths",
            dtype=offsets_dtype,
            is_input=True,
        )

        OFFSETS = ops.jagged_lengths_to_offsets()(LENGTHS)

        OFFSETS._attrs["name"] = "offsets"
        OFFSETS._attrs["is_output"] = True

        model = compile_model(
            [OFFSETS],
            detect_target(),
            "./tmp",
            f"test_jagged_lengths_to_offsets_{test_suffix}",
        )

        torch_dtype = string_to_torch_dtype(offsets_dtype)

        for seed in range(10):
            torch.manual_seed(seed)
            lengths_pt = torch.randint(
                low=0,
                high=1024,
                size=(batch_size,),
                dtype=torch_dtype,
            )
            offsets_pt = torch.cat(
                [
                    torch.zeros((1,), dtype=torch_dtype),
                    torch.cumsum(lengths_pt, dim=0, dtype=torch_dtype),
                ],
            ).cuda()

            offsets = torch.empty(
                size=(batch_size + 1,),
                dtype=torch_dtype,
            ).cuda()
            model.run_with_tensors(
                inputs={"lengths": lengths_pt.cuda()},
                outputs=[offsets],
            )

            t.assert_close(offsets, offsets_pt)
//  [python]
def _test_jagged_lengths_to_offsets(
        self,
        batch_size: int,
        offsets_dtype: str = "int32",
        test_suffix: str = "",
    ):
        LENGTHS = Tensor(
            shape=[IntVar([1, batch_size], name="batch_size")],
            name="lengths",
            dtype=offsets_dtype,
            is_input=True,
        )

        OFFSETS = ops.jagged_lengths_to_offsets()(LENGTHS)

        OFFSETS._attrs["name"] = "offsets"
        OFFSETS._attrs["is_output"] = True

        model = compile_model(
            [OFFSETS],
            detect_target(),
            "./tmp",
            f"test_jagged_lengths_to_offsets_{test_suffix}",
        )

        torch_dtype = string_to_torch_dtype(offsets_dtype)

        for seed in range(10):
            torch.manual_seed(seed)
            lengths_pt = torch.randint(
                low=0,
                high=1024,
                size=(batch_size,),
                dtype=torch_dtype,
            )
            offsets_pt = torch.cat(
                [
                    torch.zeros((1,), dtype=torch_dtype),
                    torch.cumsum(lengths_pt, dim=0, dtype=torch_dtype),
                ],
            ).cuda()

            offsets = torch.empty(
                size=(batch_size + 1,),
                dtype=torch_dtype,
            ).cuda()
            model.run_with_tensors(
                inputs={"lengths": lengths_pt.cuda()},
                outputs=[offsets],
            )

            torch.testing.assert_close(offsets, offsets_pt)
//  [python]
def test_rms_norm(self, input, weight, eps, gt):
        from lmdeploy.pytorch.kernels import rms_norm

        out = rms_norm(input, weight, eps)
        torch.testing.assert_allclose(out, gt)
//  [python]
def test_rms_norm(self, input, weight, eps, gt):
        from lmdeploy.pytorch.kernels import rms_norm

        out = rms_norm(input, weight, eps)
        torch.testing.assert_close(out, gt)
//  [python]
def test_calculate_joint_limits_exceeded(self):
        """Test that calculate_joint_limits_exceeded() is working as expected"""

        configs = torch.tensor([
            [0, 0, 0],
            [0, 0, 0],
            [-2, 0, 0],
            [0, -1.999, 0],
            [0, 2.0001, 0],
        ])
        joint_limits = [
            (-1, 1),
            (-2, 2),
            (-3, 3),
        ]
        expected = torch.tensor([False, False, True, False, True], dtype=torch.bool)
        returned = calculate_joint_limits_exceeded(configs, joint_limits)
        self.assertEqual(returned.dtype, torch.bool)
        self.assertEqual(returned.shape, (5,))
        self.assertEqual(returned, expected)
//  [python]
def test_calculate_joint_limits_exceeded(self):
        """Test that calculate_joint_limits_exceeded() is working as expected"""

        configs = torch.tensor([
            [0, 0, 0],
            [0, 0, 0],
            [-2, 0, 0],
            [0, -1.999, 0],
            [0, 2.0001, 0],
        ])
        joint_limits = [
            (-1, 1),
            (-2, 2),
            (-3, 3),
        ]
        expected = torch.tensor([False, False, True, False, True], dtype=torch.bool)
        returned = calculate_joint_limits_exceeded(configs, joint_limits)
        self.assertEqual(returned.dtype, torch.bool)
        self.assertEqual(returned.shape, (5,))
        torch.testing.assert_close(returned, expected)
//  [python]
def test_self_multihead_attn(self):
        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        torch.testing.assert_allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)

        with torch.no_grad():
            ref_grads = torch.randn_like(self.tst_inputs)
            tst_grads = ref_grads.clone()

        ref_outputs.backward(ref_grads)
        tst_outputs.backward(tst_grads)
        torch.testing.assert_allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn(self):
        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)

        with torch.no_grad():
            ref_grads = torch.randn_like(self.tst_inputs)
            tst_grads = ref_grads.clone()

        ref_outputs.backward(ref_grads)
        tst_outputs.backward(tst_grads)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn_time_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        time_mask_byte = torch.triu(
            torch.ones(
                self.tst_inputs.size(0), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        time_mask_bool = time_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=time_mask_bool,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=time_mask_byte,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        self.assertTrue(torch.allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5))
        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
//  [python]
def test_self_multihead_attn_time_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        time_mask_byte = torch.triu(
            torch.ones(
                self.tst_inputs.size(0), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        time_mask_bool = time_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=time_mask_bool,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=None,
            need_weights=False,
            attn_mask=time_mask_byte,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn_pad_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        pad_mask_byte = torch.tril(
            torch.ones(
                self.tst_inputs.size(1), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        pad_mask_bool = pad_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=pad_mask_bool,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=pad_mask_byte,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        torch.testing.assert_allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn_pad_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        pad_mask_byte = torch.tril(
            torch.ones(
                self.tst_inputs.size(1), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8
            ),
            1,
        )
        pad_mask_bool = pad_mask_byte.to(torch.bool)

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=pad_mask_bool,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=pad_mask_byte,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        data_shape = [batch_size] + list(data_shape)
        input_data = torch.rand(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand(weight_output_dim)
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_allclose(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_linear(self, batch_size, data_shape, weight_output_dim, use_bias):
        data_shape = [batch_size] + list(data_shape)
        input_data = torch.rand(data_shape)
        weight = torch.rand((weight_output_dim, data_shape[-1]))
        if use_bias:
            bias = torch.rand(weight_output_dim)
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_close(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_linear_1d_input(self, input_size, weight_output_dim, use_bias):
        input_data = torch.rand(input_size)
        weight = torch.rand((weight_output_dim, input_data.shape[-1]))
        if use_bias:
            bias = torch.rand(weight_output_dim)
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_allclose(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_linear_1d_input(self, input_size, weight_output_dim, use_bias):
        input_data = torch.rand(input_size)
        weight = torch.rand((weight_output_dim, input_data.shape[-1]))
        if use_bias:
            bias = torch.rand(weight_output_dim)
        else:
            bias = None
        ref_result = F.linear(input_data, weight, bias)
        packed_weight_bias = torch.ops.prepacked.linear_clamp_prepack(weight, bias)
        output_linearprepacked = torch.ops.prepacked.linear_clamp_run(input_data, packed_weight_bias)
        torch.testing.assert_close(ref_result, output_linearprepacked, rtol=1e-2, atol=1e-3)
//  [python]
def test_auc(self, n_tasks: int, batch_size: int, dtype: torch.dtype) -> None:
        predictions = torch.randint(0, 1000, (n_tasks, batch_size)).to(dtype).cuda()
        labels = torch.randint(0, 1000, (n_tasks, batch_size)).to(dtype).cuda() / 1000.0
        weights = torch.rand(n_tasks, batch_size).to(dtype).cuda()

        compute_auc = fbgemm_gpu.metrics.Auc()
        output_ref = compute_auc(n_tasks, predictions, labels, weights)
        output = fbgemm_gpu.metrics.auc(n_tasks, predictions, labels, weights)

        # Explicitly convert type based on output_ref's dtype
        output = output.to(output_ref.dtype)

        # Test correctness only if output_ref does not product nan or inf
        if not (torch.isnan(output_ref).any() or torch.isinf(output_ref).any()):
            t.assert_close(
                output_ref,
                output,
                rtol=1e-2 if dtype == torch.half else None,
                atol=1e-2 if dtype == torch.half else None,
            )
//  [python]
def test_auc(self, n_tasks: int, batch_size: int, dtype: torch.dtype) -> None:
        predictions = torch.randint(0, 1000, (n_tasks, batch_size)).to(dtype).cuda()
        labels = torch.randint(0, 1000, (n_tasks, batch_size)).to(dtype).cuda() / 1000.0
        weights = torch.rand(n_tasks, batch_size).to(dtype).cuda()

        compute_auc = fbgemm_gpu.metrics.Auc()
        output_ref = compute_auc(n_tasks, predictions, labels, weights)
        output = fbgemm_gpu.metrics.auc(n_tasks, predictions, labels, weights)

        # Explicitly convert type based on output_ref's dtype
        output = output.to(output_ref.dtype)

        # Test correctness only if output_ref does not product nan or inf
        if not (torch.isnan(output_ref).any() or torch.isinf(output_ref).any()):
            torch.testing.assert_close(
                output_ref,
                output,
                rtol=1e-2 if dtype == torch.half else None,
                atol=1e-2 if dtype == torch.half else None,
            )
//  [python]
def test_incremental_write(tmp_path):
    sd = {str(k): torch.randn(5, 10) for k in range(3)}
    sd["0"].someattr = 1
    sd_expected = {k: v.clone() for k, v in sd.items()}
    fn = str(tmp_path / "test.pt")
    with incremental_save(fn) as f:
        sd["0"] = f.store_early(sd["0"])
        sd["2"] = f.store_early(sd["2"])
        f.save(sd)
    sd_actual = torch.load(fn)
    assert sd_actual.keys() == sd_expected.keys()
    assert sd_actual["0"].someattr == 1  # requires PyTorch 2.0+
    for k, v_expected in sd_expected.items():
        v_actual = sd_actual[k]
        torch.testing.assert_allclose(v_expected, v_actual)
//  [python]
def test_incremental_write(tmp_path):
    sd = {str(k): torch.randn(5, 10) for k in range(3)}
    sd["0"].someattr = 1
    sd_expected = {k: v.clone() for k, v in sd.items()}
    fn = str(tmp_path / "test.pt")
    with incremental_save(fn) as f:
        sd["0"] = f.store_early(sd["0"])
        sd["2"] = f.store_early(sd["2"])
        f.save(sd)
    sd_actual = torch.load(fn)
    assert sd_actual.keys() == sd_expected.keys()
    assert sd_actual["0"].someattr == 1  # requires PyTorch 2.0+
    for k, v_expected in sd_expected.items():
        v_actual = sd_actual[k]
        torch.testing.assert_close(v_expected, v_actual)
//  [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().__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().__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().__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().__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 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)
            t.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)
            t.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)