You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/11/24 01:18:17 UTC
[GitHub] eric-haibin-lin closed pull request #13290: A few operators on
graphs stored as CSR
eric-haibin-lin closed pull request #13290: A few operators on graphs stored as CSR
URL: https://github.com/apache/incubator-mxnet/pull/13290
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/docs/api/python/ndarray/contrib.md b/docs/api/python/ndarray/contrib.md
index b2bfa170b1d..709ddae007c 100644
--- a/docs/api/python/ndarray/contrib.md
+++ b/docs/api/python/ndarray/contrib.md
@@ -60,6 +60,7 @@ In the rest of this document, we list routines provided by the `ndarray.contrib`
isnan
index_copy
getnnz
+ edge_id
```
## API Reference
diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc
new file mode 100644
index 00000000000..bf54ed398f9
--- /dev/null
+++ b/src/operator/contrib/dgl_graph.cc
@@ -0,0 +1,464 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <mxnet/io.h>
+#include <mxnet/base.h>
+#include <mxnet/ndarray.h>
+#include <mxnet/operator.h>
+#include <mxnet/operator_util.h>
+#include <dmlc/logging.h>
+#include <dmlc/optional.h>
+#include "../operator_common.h"
+#include "../elemwise_op_common.h"
+#include "../../imperative/imperative_utils.h"
+#include "../subgraph_op_common.h"
+#include "../mshadow_op.h"
+#include "../mxnet_op.h"
+#include "../tensor/init_op.h"
+
+namespace mxnet {
+namespace op {
+
+
+///////////////////////// Create induced subgraph ///////////////////////////
+
+struct DGLSubgraphParam : public dmlc::Parameter<DGLSubgraphParam> {
+ int num_args;
+ bool return_mapping;
+ DMLC_DECLARE_PARAMETER(DGLSubgraphParam) {
+ DMLC_DECLARE_FIELD(num_args).set_lower_bound(2)
+ .describe("Number of input arguments, including all symbol inputs.");
+ DMLC_DECLARE_FIELD(return_mapping)
+ .describe("Return mapping of vid and eid between the subgraph and the parent graph.");
+ }
+}; // struct DGLSubgraphParam
+
+DMLC_REGISTER_PARAMETER(DGLSubgraphParam);
+
+static bool DGLSubgraphStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->at(0), kCSRStorage);
+ for (size_t i = 1; i < in_attrs->size(); i++)
+ CHECK_EQ(in_attrs->at(i), kDefaultStorage);
+
+ bool success = true;
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ if (!type_assign(&(*out_attrs)[i], mxnet::kCSRStorage))
+ success = false;
+ }
+ return success;
+}
+
+static bool DGLSubgraphShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape> *in_attrs,
+ std::vector<TShape> *out_attrs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->at(0).ndim(), 2U);
+ for (size_t i = 1; i < in_attrs->size(); i++)
+ CHECK_EQ(in_attrs->at(i).ndim(), 1U);
+
+ size_t num_g = params.num_args - 1;
+ for (size_t i = 0; i < num_g; i++) {
+ TShape gshape(2);
+ gshape[0] = in_attrs->at(i + 1)[0];
+ gshape[1] = in_attrs->at(i + 1)[0];
+ out_attrs->at(i) = gshape;
+ }
+ for (size_t i = num_g; i < out_attrs->size(); i++) {
+ TShape gshape(2);
+ gshape[0] = in_attrs->at(i - num_g + 1)[0];
+ gshape[1] = in_attrs->at(i - num_g + 1)[0];
+ out_attrs->at(i) = gshape;
+ }
+ return true;
+}
+
+static bool DGLSubgraphType(const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ size_t num_g = params.num_args - 1;
+ for (size_t i = 0; i < num_g; i++) {
+ CHECK_EQ(in_attrs->at(i + 1), mshadow::kInt64);
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ out_attrs->at(i) = in_attrs->at(0);
+ }
+ return true;
+}
+
+typedef int64_t dgl_id_t;
+
+class Bitmap {
+ const size_t size = 1024 * 1024 * 4;
+ const size_t mask = size - 1;
+ std::vector<bool> map;
+
+ size_t hash(dgl_id_t id) const {
+ return id & mask;
+ }
+ public:
+ Bitmap(const dgl_id_t *vid_data, int64_t len): map(size) {
+ for (int64_t i = 0; i < len; ++i) {
+ map[hash(vid_data[i])] = 1;
+ }
+ }
+
+ bool test(dgl_id_t id) const {
+ return map[hash(id)];
+ }
+};
+
+/*
+ * This uses a hashtable to check if a node is in the given node list.
+ */
+class HashTableChecker {
+ std::unordered_map<dgl_id_t, dgl_id_t> oldv2newv;
+ Bitmap map;
+
+ public:
+ HashTableChecker(const dgl_id_t *vid_data, int64_t len): map(vid_data, len) {
+ oldv2newv.reserve(len);
+ for (int64_t i = 0; i < len; ++i) {
+ oldv2newv[vid_data[i]] = i;
+ }
+ }
+
+ void CollectOnRow(const dgl_id_t col_idx[], const dgl_id_t eids[], size_t row_len,
+ std::vector<dgl_id_t> *new_col_idx,
+ std::vector<dgl_id_t> *orig_eids) {
+ // TODO(zhengda) I need to make sure the column index in each row is sorted.
+ for (size_t j = 0; j < row_len; ++j) {
+ const dgl_id_t oldsucc = col_idx[j];
+ const dgl_id_t eid = eids[j];
+ Collect(oldsucc, eid, new_col_idx, orig_eids);
+ }
+ }
+
+ void Collect(const dgl_id_t old_id, const dgl_id_t old_eid,
+ std::vector<dgl_id_t> *col_idx,
+ std::vector<dgl_id_t> *orig_eids) {
+ if (!map.test(old_id))
+ return;
+
+ auto it = oldv2newv.find(old_id);
+ if (it != oldv2newv.end()) {
+ const dgl_id_t new_id = it->second;
+ col_idx->push_back(new_id);
+ if (orig_eids)
+ orig_eids->push_back(old_eid);
+ }
+ }
+};
+
+static void GetSubgraph(const NDArray &csr_arr, const NDArray &varr,
+ const NDArray &sub_csr, const NDArray *old_eids) {
+ const TBlob &data = varr.data();
+ int64_t num_vertices = csr_arr.shape()[0];
+ const size_t len = varr.shape()[0];
+ const dgl_id_t *vid_data = data.dptr<dgl_id_t>();
+ HashTableChecker def_check(vid_data, len);
+ // check if varr is sorted.
+ CHECK(std::is_sorted(vid_data, vid_data + len)) << "The input vertex list has to be sorted";
+
+ // Collect the non-zero entries in from the original graph.
+ std::vector<dgl_id_t> row_idx(len + 1);
+ std::vector<dgl_id_t> col_idx;
+ std::vector<dgl_id_t> orig_eids;
+ col_idx.reserve(len * 50);
+ orig_eids.reserve(len * 50);
+ const dgl_id_t *eids = csr_arr.data().dptr<dgl_id_t>();
+ const dgl_id_t *indptr = csr_arr.aux_data(csr::kIndPtr).dptr<dgl_id_t>();
+ const dgl_id_t *indices = csr_arr.aux_data(csr::kIdx).dptr<dgl_id_t>();
+ for (size_t i = 0; i < len; ++i) {
+ const dgl_id_t oldvid = vid_data[i];
+ CHECK_LT(oldvid, num_vertices) << "Vertex Id " << oldvid << " isn't in a graph of "
+ << num_vertices << " vertices";
+ size_t row_start = indptr[oldvid];
+ size_t row_len = indptr[oldvid + 1] - indptr[oldvid];
+ def_check.CollectOnRow(indices + row_start, eids + row_start, row_len,
+ &col_idx, old_eids == nullptr ? nullptr : &orig_eids);
+
+ row_idx[i + 1] = col_idx.size();
+ }
+
+ TShape nz_shape(1);
+ nz_shape[0] = col_idx.size();
+ TShape indptr_shape(1);
+ indptr_shape[0] = row_idx.size();
+
+ // Store the non-zeros in a subgraph with edge attributes of new edge ids.
+ sub_csr.CheckAndAllocData(nz_shape);
+ sub_csr.CheckAndAllocAuxData(csr::kIdx, nz_shape);
+ sub_csr.CheckAndAllocAuxData(csr::kIndPtr, indptr_shape);
+ dgl_id_t *indices_out = sub_csr.aux_data(csr::kIdx).dptr<dgl_id_t>();
+ dgl_id_t *indptr_out = sub_csr.aux_data(csr::kIndPtr).dptr<dgl_id_t>();
+ std::copy(col_idx.begin(), col_idx.end(), indices_out);
+ std::copy(row_idx.begin(), row_idx.end(), indptr_out);
+ dgl_id_t *sub_eids = sub_csr.data().dptr<dgl_id_t>();
+ for (int64_t i = 0; i < nz_shape[0]; i++)
+ sub_eids[i] = i;
+
+ // Store the non-zeros in a subgraph with edge attributes of old edge ids.
+ if (old_eids) {
+ old_eids->CheckAndAllocData(nz_shape);
+ old_eids->CheckAndAllocAuxData(csr::kIdx, nz_shape);
+ old_eids->CheckAndAllocAuxData(csr::kIndPtr, indptr_shape);
+ dgl_id_t *indices_out = old_eids->aux_data(csr::kIdx).dptr<dgl_id_t>();
+ dgl_id_t *indptr_out = old_eids->aux_data(csr::kIndPtr).dptr<dgl_id_t>();
+ dgl_id_t *sub_eids = old_eids->data().dptr<dgl_id_t>();
+ std::copy(col_idx.begin(), col_idx.end(), indices_out);
+ std::copy(row_idx.begin(), row_idx.end(), indptr_out);
+ std::copy(orig_eids.begin(), orig_eids.end(), sub_eids);
+ }
+}
+
+static void DGLSubgraphComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ int num_g = params.num_args - 1;
+#pragma omp parallel for
+ for (int i = 0; i < num_g; i++) {
+ const NDArray *old_eids = params.return_mapping ? &outputs[i + num_g] : nullptr;
+ GetSubgraph(inputs[0], inputs[i + 1], outputs[i], old_eids);
+ }
+}
+
+NNVM_REGISTER_OP(_contrib_dgl_subgraph)
+.describe(R"code(This operator constructs an induced subgraph for
+a given set of vertices from a graph. The operator accepts multiple
+sets of vertices as input. For each set of vertices, it returns a pair
+of CSR matrices if return_mapping is True: the first matrix contains edges
+with new edge Ids, the second matrix contains edges with the original
+edge Ids.
+Example::
+ x=[[1, 0, 0, 2],
+ [3, 0, 4, 0],
+ [0, 5, 0, 0],
+ [0, 6, 7, 0]]
+ v = [0, 1, 2]
+ dgl_subgraph(x, v, return_mapping=True) =
+ [[1, 0, 0],
+ [2, 0, 3],
+ [0, 4, 0]],
+ [[1, 0, 0],
+ [3, 0, 4],
+ [0, 5, 0]]
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<DGLSubgraphParam>)
+.set_num_inputs([](const NodeAttrs& attrs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ return params.num_args;
+})
+.set_num_outputs([](const NodeAttrs& attrs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ int num_varray = params.num_args - 1;
+ if (params.return_mapping)
+ return num_varray * 2;
+ else
+ return num_varray;
+})
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ const DGLSubgraphParam& params = nnvm::get<DGLSubgraphParam>(attrs.parsed);
+ std::vector<std::string> names;
+ names.reserve(params.num_args);
+ names.emplace_back("graph");
+ for (int i = 1; i < params.num_args; ++i)
+ names.push_back("varray" + std::to_string(i - 1));
+ return names;
+})
+.set_attr<FInferStorageType>("FInferStorageType", DGLSubgraphStorageType)
+.set_attr<nnvm::FInferShape>("FInferShape", DGLSubgraphShape)
+.set_attr<nnvm::FInferType>("FInferType", DGLSubgraphType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", DGLSubgraphComputeExCPU)
+.set_attr<std::string>("key_var_num_args", "num_args")
+.add_argument("graph", "NDArray-or-Symbol", "Input graph where we sample vertices.")
+.add_argument("data", "NDArray-or-Symbol[]",
+ "The input arrays that include data arrays and states.")
+.add_arguments(DGLSubgraphParam::__FIELDS__());
+
+///////////////////////// Edge Id ///////////////////////////
+
+inline bool EdgeIDShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape>* in_attrs,
+ std::vector<TShape>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 3U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ CHECK_EQ(in_attrs->at(1).ndim(), 1U);
+ CHECK_EQ(in_attrs->at(2).ndim(), 1U);
+ CHECK_EQ(in_attrs->at(1)[0], in_attrs->at(2)[0]);
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(1));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 1, out_attrs->at(0));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 2, out_attrs->at(0));
+ return out_attrs->at(0).ndim() != 0U && out_attrs->at(0).Size() != 0U;
+}
+
+inline bool EdgeIDType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 3U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
+ return out_attrs->at(0) != -1;
+}
+
+inline bool EdgeIDStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 3U) << "Only works for 2d arrays";
+ CHECK_EQ(out_attrs->size(), 1U);
+ int& in_stype = in_attrs->at(0);
+ int& out_stype = out_attrs->at(0);
+ bool dispatched = false;
+ if (!dispatched && in_stype == kCSRStorage) {
+ // csr -> dns
+ dispatched = storage_type_assign(&out_stype, kDefaultStorage,
+ dispatch_mode, DispatchMode::kFComputeEx);
+ }
+ if (!dispatched) {
+ LOG(ERROR) << "Cannot dispatch edge_id storage type, only works for csr matrices";
+ }
+ return dispatched;
+}
+
+struct edge_id_csr_forward {
+ template<typename DType, typename IType, typename CType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
+ const IType* in_indices, const IType* in_indptr,
+ const CType* u, const CType* v) {
+ const int64_t target_row_id = static_cast<int64_t>(u[i]);
+ const IType target_col_id = static_cast<IType>(v[i]);
+ auto ptr = std::find(in_indices + in_indptr[target_row_id],
+ in_indices + in_indptr[target_row_id + 1], target_col_id);
+ if (ptr == in_indices + in_indptr[target_row_id + 1]) {
+ // does not exist in the range
+ out_data[i] = DType(-1);
+ } else {
+ out_data[i] = *(in_data + (ptr - in_indices));
+ }
+ }
+};
+
+template<typename xpu>
+void EdgeIDForwardCsrImpl(const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const OpReqType req,
+ const NDArray& output) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using namespace csr;
+ if (req == kNullOp) return;
+ CHECK_EQ(inputs.size(), 3U);
+ CHECK_EQ(req, kWriteTo) << "EdgeID with CSR only supports kWriteTo";
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ const NDArray& u = inputs[1];
+ const nnvm::dim_t out_elems = u.shape().Size();
+ if (!inputs[0].storage_initialized()) {
+ MSHADOW_TYPE_SWITCH(output.dtype(), DType, {
+ Kernel<mxnet_op::op_with_req<mshadow_op::identity, kWriteTo>, xpu>::Launch(
+ s, out_elems, output.data().dptr<DType>(), DType(-1));
+ });
+ return;
+ }
+ const NDArray& data = inputs[0];
+ const TBlob& in_data = data.data();
+ const TBlob& in_indices = data.aux_data(kIdx);
+ const TBlob& in_indptr = data.aux_data(kIndPtr);
+ const NDArray& v = inputs[2];
+
+ CHECK_EQ(data.aux_type(kIdx), data.aux_type(kIndPtr))
+ << "The dtypes of indices and indptr don't match";
+ MSHADOW_TYPE_SWITCH(data.dtype(), DType, {
+ MSHADOW_IDX_TYPE_SWITCH(data.aux_type(kIdx), IType, {
+ MSHADOW_TYPE_SWITCH(u.dtype(), CType, {
+ Kernel<edge_id_csr_forward, xpu>::Launch(
+ s, out_elems, output.data().dptr<DType>(), in_data.dptr<DType>(),
+ in_indices.dptr<IType>(), in_indptr.dptr<IType>(),
+ u.data().dptr<CType>(), v.data().dptr<CType>());
+ });
+ });
+ });
+}
+
+template<typename xpu>
+void EdgeIDForwardEx(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ CHECK_EQ(inputs.size(), 3U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ const auto in_stype = inputs[0].storage_type();
+ const auto out_stype = outputs[0].storage_type();
+ if (in_stype == kCSRStorage && out_stype == kDefaultStorage) {
+ EdgeIDForwardCsrImpl<xpu>(ctx, inputs, req[0], outputs[0]);
+ } else {
+ LogUnimplementedOp(attrs, ctx, inputs, req, outputs);
+ }
+}
+
+NNVM_REGISTER_OP(_contrib_edge_id)
+.describe(R"code(This operator implements the edge_id function for a graph
+stored in a CSR matrix (the value of the CSR stores the edge Id of the graph).
+output[i] = input[u[i], v[i]] if there is an edge between u[i] and v[i]],
+otherwise output[i] will be -1. Both u and v should be 1D vectors.
+Example::
+ x = [[ 1, 0, 0 ],
+ [ 0, 2, 0 ],
+ [ 0, 0, 3 ]]
+ u = [ 0, 0, 1, 1, 2, 2 ]
+ v = [ 0, 1, 1, 2, 0, 2 ]
+ edge_id(x, u, v) = [ 1, -1, 2, -1, -1, 3 ]
+
+The storage type of ``edge_id`` output depends on storage types of inputs
+ - edge_id(csr, default, default) = default
+ - default and rsp inputs are not supported
+
+)code" ADD_FILELINE)
+.set_num_inputs(3)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data", "u", "v"};
+ })
+.set_attr<nnvm::FInferShape>("FInferShape", EdgeIDShape)
+.set_attr<nnvm::FInferType>("FInferType", EdgeIDType)
+.set_attr<FInferStorageType>("FInferStorageType", EdgeIDStorageType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", EdgeIDForwardEx<cpu>)
+.add_argument("data", "NDArray-or-Symbol", "Input ndarray")
+.add_argument("u", "NDArray-or-Symbol", "u ndarray")
+.add_argument("v", "NDArray-or-Symbol", "v ndarray");
+
+
+} // namespace op
+} // namespace mxnet
diff --git a/tests/python/unittest/test_contrib_operator.py b/tests/python/unittest/test_contrib_operator.py
index 58728d8dad6..43d3db648a8 100644
--- a/tests/python/unittest/test_contrib_operator.py
+++ b/tests/python/unittest/test_contrib_operator.py
@@ -261,6 +261,7 @@ def test_multibox_target_op():
assert_array_equal(loc_mask.asnumpy(), expected_loc_mask)
assert_array_equal(cls_target.asnumpy(), expected_cls_target)
+
if __name__ == '__main__':
import nose
nose.runmodule()
diff --git a/tests/python/unittest/test_dgl_graph.py b/tests/python/unittest/test_dgl_graph.py
new file mode 100644
index 00000000000..774f811a0d1
--- /dev/null
+++ b/tests/python/unittest/test_dgl_graph.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: skip-file
+from __future__ import print_function
+import numpy as np
+import scipy as sp
+import mxnet as mx
+import random
+import itertools
+from numpy.testing import assert_allclose, assert_array_equal
+from mxnet.test_utils import *
+import unittest
+
+def test_edge_id():
+ shape = rand_shape_2d()
+ data = rand_ndarray(shape, stype='csr', density=0.4)
+ ground_truth = np.zeros(shape, dtype=np.float32)
+ ground_truth -= 1.0
+ indptr_np = data.indptr.asnumpy()
+ data_np = data.data.asnumpy()
+ indices_np = data.indices.asnumpy()
+ for i in range(shape[0]):
+ for j in range(indptr_np[i], indptr_np[i+1]):
+ idx = indices_np[j]
+ ground_truth[i, idx] = data_np[j]
+
+ np_u = np.random.randint(0, shape[0], size=(5, ))
+ np_v = np.random.randint(0, shape[1], size=(5, ))
+ mx_u = mx.nd.array(np_u)
+ mx_v = mx.nd.array(np_v)
+ assert_almost_equal(mx.nd.contrib.edge_id(data, mx_u, mx_v).asnumpy(),
+ ground_truth[np_u, np_v], rtol=1e-5, atol=1e-6)
+
+def generate_graph(n):
+ arr = sp.sparse.random(n, n, density=0.2, format='coo')
+ arr.data = np.arange(0, len(arr.row), dtype=np.float32)
+ return arr.tocsr(), mx.nd.sparse.csr_matrix(arr.tocsr()).astype(np.int64)
+
+def test_subgraph():
+ sp_g, g = generate_graph(100)
+ vertices = np.unique(np.random.randint(0, 100, size=(20)))
+ subgs = mx.nd.contrib.dgl_subgraph(g, mx.nd.array(vertices, dtype=np.int64),
+ return_mapping=True)
+ subgs[0].check_format()
+ subgs[1].check_format()
+ assert_array_equal(subgs[0].indptr, subgs[1].indptr)
+ assert_array_equal(subgs[0].indices, subgs[1].indices)
+ sp_subg = subgs[1].asscipy()
+ for i in range(len(subgs[0].indptr) - 1):
+ subv1 = i
+ v1 = vertices[subv1]
+ row_start = int(subgs[0].indptr[subv1].asnumpy()[0])
+ row_end = int(subgs[0].indptr[subv1 + 1].asnumpy()[0])
+ if row_start >= len(subgs[0].indices):
+ remain = subgs[0].indptr[subv1:].asnumpy()
+ assert np.sum(remain == row_start) == len(remain)
+ break
+ row = subgs[0].indices[row_start:row_end]
+ for j, subv2 in enumerate(row.asnumpy()):
+ v2 = vertices[subv2]
+ assert sp_g[v1, v2] == sp_subg[subv1, subv2]
+
+if __name__ == "__main__":
+ import nose
+ nose.runmodule()
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services