This is an automated email from the ASF dual-hosted git repository.
haibin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new a7d17a4 add subgraph adjacency operator. (#13396)
a7d17a4 is described below
commit a7d17a42a08e6006605533a6dc2a3873040efca1
Author: BullDemonKing <[email protected]>
AuthorDate: Thu Nov 29 03:20:47 2018 +0800
add subgraph adjacency operator. (#13396)
* add adjacency.
* fix lint.
* add GPU impl.
* retrigger
* address comments.
* Update dgl_graph.cc
---
src/operator/contrib/dgl_graph-inl.h | 68 +++++++++++++++++++++++++++++++
src/operator/contrib/dgl_graph.cc | 72 +++++++++++++++++++++++++++++++--
src/operator/contrib/dgl_graph.cu | 29 +++++++++++++
tests/python/unittest/test_dgl_graph.py | 9 +++++
4 files changed, 174 insertions(+), 4 deletions(-)
diff --git a/src/operator/contrib/dgl_graph-inl.h
b/src/operator/contrib/dgl_graph-inl.h
new file mode 100644
index 0000000..f31071b
--- /dev/null
+++ b/src/operator/contrib/dgl_graph-inl.h
@@ -0,0 +1,68 @@
+/*
+ * 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.
+ */
+
+#ifndef MXNET_OPERATOR_CONTRIB_DGL_GRAPH_INL_H_
+#define MXNET_OPERATOR_CONTRIB_DGL_GRAPH_INL_H_
+
+#include <dmlc/logging.h>
+#include <dmlc/parameter.h>
+#include <mxnet/operator.h>
+#include <mxnet/ndarray.h>
+#include <map>
+#include <algorithm>
+#include <vector>
+#include <string>
+#include <utility>
+#include "../operator_common.h"
+#include "../mxnet_op.h"
+#include "../mshadow_op.h"
+#include "../tensor/init_op.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename xpu>
+void DGLAdjacencyForwardEx(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(), 1U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ CHECK_EQ(inputs[0].storage_type(), kCSRStorage);
+ CHECK_EQ(outputs[0].storage_type(), kCSRStorage);
+ CHECK_EQ(req[0], kWriteTo);
+ const TBlob &in_idx = inputs[0].aux_data(csr::kIdx);
+ const TBlob &in_indptr = inputs[0].aux_data(csr::kIndPtr);
+
+ outputs[0].CheckAndAllocData(in_idx.shape_);
+ outputs[0].CheckAndAllocAuxData(csr::kIdx, in_idx.shape_);
+ outputs[0].CheckAndAllocAuxData(csr::kIndPtr, in_indptr.shape_);
+
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ Fill<false>(s, outputs[0].data(), req[0], 1.0);
+ mxnet_op::copy(s, outputs[0].aux_data(csr::kIdx), in_idx);
+ mxnet_op::copy(s, outputs[0].aux_data(csr::kIndPtr), in_indptr);
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_OPERATOR_CONTRIB_DGL_GRAPH_INL_H_
diff --git a/src/operator/contrib/dgl_graph.cc
b/src/operator/contrib/dgl_graph.cc
index bf54ed3..d9bcdd4 100644
--- a/src/operator/contrib/dgl_graph.cc
+++ b/src/operator/contrib/dgl_graph.cc
@@ -24,13 +24,10 @@
#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"
+#include "./dgl_graph-inl.h"
namespace mxnet {
namespace op {
@@ -459,6 +456,73 @@ The storage type of ``edge_id`` output depends on storage
types of inputs
.add_argument("u", "NDArray-or-Symbol", "u ndarray")
.add_argument("v", "NDArray-or-Symbol", "v ndarray");
+///////////////////////// DGL Adjacency ///////////////////////////
+
+inline bool DGLAdjacencyShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape>* in_attrs,
+ std::vector<TShape>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
+ return out_attrs->at(0).ndim() != 0U && out_attrs->at(0).Size() != 0U;
+}
+
+inline bool DGLAdjacencyType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ CHECK_EQ(in_attrs->at(0), mshadow::kInt64);
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kFloat32);
+ return out_attrs->at(0) != -1;
+}
+
+inline bool DGLAdjacencyStorageType(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(), 1U) << "Only works for 2d arrays";
+ CHECK_EQ(out_attrs->size(), 1U);
+ int& out_stype = out_attrs->at(0);
+ bool dispatched = storage_type_assign(&out_stype, kCSRStorage,
+ dispatch_mode,
DispatchMode::kFComputeEx);
+ if (!dispatched) {
+ LOG(ERROR) << "Cannot dispatch output storage type: " <<
common::stype_string(out_stype)
+ << ". dgl_adjacency only works for csr matrices";
+ }
+ return dispatched;
+}
+
+NNVM_REGISTER_OP(_contrib_dgl_adjacency)
+.describe(R"code(This operator converts a CSR matrix whose values are edge Ids
+to an adjacency matrix whose values are ones. The output CSR matrix always has
+the data value of float32.
+Example::
+
+ x = [[ 1, 0, 0 ],
+ [ 0, 2, 0 ],
+ [ 0, 0, 3 ]]
+ dgl_adjacency(x) =
+ [[ 1, 0, 0 ],
+ [ 0, 1, 0 ],
+ [ 0, 0, 1 ]]
+
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data"};
+ })
+.set_attr<nnvm::FInferShape>("FInferShape", DGLAdjacencyShape)
+.set_attr<nnvm::FInferType>("FInferType", DGLAdjacencyType)
+.set_attr<FInferStorageType>("FInferStorageType", DGLAdjacencyStorageType)
+.set_attr<FComputeEx>("FComputeEx<cpu>", DGLAdjacencyForwardEx<cpu>)
+.add_argument("data", "NDArray-or-Symbol", "Input ndarray");
+
} // namespace op
} // namespace mxnet
diff --git a/src/operator/contrib/dgl_graph.cu
b/src/operator/contrib/dgl_graph.cu
new file mode 100644
index 0000000..336e9d4
--- /dev/null
+++ b/src/operator/contrib/dgl_graph.cu
@@ -0,0 +1,29 @@
+/*
+ * 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 "dgl_graph-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_contrib_dgl_adjacency)
+.set_attr<FComputeEx>("FComputeEx<gpu>", DGLAdjacencyForwardEx<gpu>);
+
+} // namespace op
+} // namespace mxnet
diff --git a/tests/python/unittest/test_dgl_graph.py
b/tests/python/unittest/test_dgl_graph.py
index 774f811..1d758b9 100644
--- a/tests/python/unittest/test_dgl_graph.py
+++ b/tests/python/unittest/test_dgl_graph.py
@@ -75,6 +75,15 @@ def test_subgraph():
v2 = vertices[subv2]
assert sp_g[v1, v2] == sp_subg[subv1, subv2]
+def test_adjacency():
+ sp_g, g = generate_graph(100)
+ adj = mx.nd.contrib.dgl_adjacency(g)
+ assert adj.dtype == np.float32
+ assert adj.shape == g.shape
+ assert_array_equal(adj.indptr, g.indptr)
+ assert_array_equal(adj.indices, g.indices)
+ assert_array_equal(adj.data, mx.nd.ones(shape=g.indices.shape))
+
if __name__ == "__main__":
import nose
nose.runmodule()
