diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 89ba4fc4b40d..ed7caacfdbae 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -768,7 +768,10 @@ static void CSRNeighborUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs, NNVM_REGISTER_OP(_contrib_dgl_csr_neighbor_uniform_sample) .describe(R"code(This operator samples sub-graph from a csr graph via an uniform probability. -Example:: + +Example: + + .. code:: python shape = (5, 5) data_np = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], dtype=np.int64) @@ -850,7 +853,10 @@ static void CSRNeighborNonUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs NNVM_REGISTER_OP(_contrib_dgl_csr_neighbor_non_uniform_sample) .describe(R"code(This operator samples sub-graph from a csr graph via an uniform probability. -Example:: + +Example: + + .. code:: python shape = (5, 5) prob = mx.nd.array([0.9, 0.8, 0.2, 0.4, 0.1], dtype=np.float32) @@ -1379,6 +1385,8 @@ the data value of float32. Example: + .. code:: python + x = [[ 1, 0, 0 ], [ 0, 2, 0 ], [ 0, 0, 3 ]] @@ -1549,14 +1557,33 @@ NNVM_REGISTER_OP(_contrib_dgl_graph_compact) .describe(R"code(This operator compacts a CSR matrix generated by dgl_csr_neighbor_uniform_sample and dgl_csr_neighbor_non_uniform_sample. The CSR matrices generated by these two operators may have many empty -rows at the end. This operator removes these empty rows and empty columns. -Example:: - subgs = mx.nd.contrib.csr_neighbor_uniform_sample(csr, seed, num_hops=1, - num_neighbor=2, max_num_vertices=5) - subg_v = subgs[0] - subg = subgs[1] - compacts = mx.nd.contrib.dgl_graph_compact(subg, subg_v, - graph_sizes=(subg_v[-1].asnumpy()[0])) +rows at the end and many empty columns. This operator removes these +empty rows and empty columns. + +Example: + + .. code:: python + + shape = (5, 5) + data_np = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], dtype=np.int64) + indices_np = np.array([1,2,3,4,0,2,3,4,0,1,3,4,0,1,2,4,0,1,2,3], dtype=np.int64) + indptr_np = np.array([0,4,8,12,16,20], dtype=np.int64) + a = mx.nd.sparse.csr_matrix((data_np, indices_np, indptr_np), shape=shape) + seed = mx.nd.array([0,1,2,3,4], dtype=np.int64) + out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, + num_neighbor=2, max_num_vertices=6) + subg_v = out[0] + subg = out[1] + compact = mx.nd.contrib.dgl_graph_compact(subg, subg_v, + graph_sizes=(subg_v[-1].asnumpy()[0]), return_mapping=False) + + compact.asnumpy() + array([[0, 0, 0, 1, 0], + [2, 0, 3, 0, 0], + [0, 4, 0, 0, 5], + [0, 6, 0, 0, 7], + [8, 9, 0, 0, 0]]) + )code" ADD_FILELINE) .set_attr_parser(ParamParser) .set_num_inputs([](const NodeAttrs& attrs) {