WIP: topdown example

6858ed54 · Minjie Wang · 5b5baee1 · 6858ed54
Commit 6858ed54 authored Jun 15, 2018 by Minjie Wang
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 3 deletions

python/dgl/graph.py python/dgl/graph.py +8 -3

No files found.
--- a/python/dgl/graph.py
+++ b/python/dgl/graph.py
@@ -244,9 +244,14 @@ class DGLGraph(DiGraph):
        preds = defaultdict(list)
        for uu, vv in utils.edge_iter(u, v):
            preds[vv].append(uu)
-        dst = preds.keys()
+        if len(preds) == 1:
-        src = [preds[d] for d in dst]
+            dst = list(preds.keys())[0]
-        self.recvfrom(dst, src)
+            src = preds[dst]
+            self.recvfrom(dst, src)
+        elif len(preds) > 1:
+            dst = list(preds.keys())
+            src = [preds[d] for d in dst]
+            self.recvfrom(dst, src)
    def update_to(self, u):
        """Pull messages from the node's predecessors and then update it.