Function _random_simple_graph

tests/python/common/test_sparse_ops-csr.py:15–41

(
    idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype
)

def _random_simple_graph(
    idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype
):
    src = np.random.randint(0, M, (max_nnz,))
    dst = np.random.randint(0, N, (max_nnz,))
    val = np.random.randn(max_nnz)
    a = ssp.csr_matrix((val, (src, dst)), shape=(M, N))
    a.sum_duplicates()
    a = a.tocoo()
    # shuffle edges
    perm = np.random.permutation(a.nnz)
    row = a.row[perm]
    col = a.col[perm]
    val = a.data[perm]
    a = ssp.csr_matrix((val, (row, col)), shape=(M, N))

    A = dgl.heterograph(
        {
            (srctype, etype, dsttype): (
                F.copy_to(F.tensor(row, dtype=idtype), ctx),
                F.copy_to(F.tensor(col, dtype=idtype), ctx),
            )
        },
        num_nodes_dict={srctype: a.shape[0], dsttype: a.shape[1]},
    )
    A.edata["w"] = F.copy_to(F.tensor(val, dtype=dtype), ctx)
    return a, A


@parametrize_idtype