refine

dmlc · May 24, 2024 · 714b62b · 714b62b
1 parent e206643
commit 714b62b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/examples/distributed/graphsage/node_classification_unsupervised.py b/examples/distributed/graphsage/node_classification_unsupervised.py
@@ -269,7 +269,9 @@ def run(args, device, data):
                         elif exclude == "reverse_id":
                             assert not th.any(th.isin(current_eids, seed_eids))
                         else:
-                            raise ValueError(f"Unsupported exclude type: {exclude}")
+                            raise ValueError(
+                                f"Unsupported exclude type: {exclude}"
+                            )
                 tic_step = time.time()
                 sample_t.append(tic_step - start)
 

diff --git a/python/dgl/distributed/graph_services.py b/python/dgl/distributed/graph_services.py
@@ -956,7 +956,7 @@ def issue_remote_req(node_ids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
-            exclude_edges=None,
+            exclude_edges=exclude_edges,
             replace=replace,
             etype_sorted=etype_sorted,
             use_graphbolt=use_graphbolt,
@@ -984,7 +984,7 @@ def local_access(local_g, partition_book, local_nids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
-            exclude_edges=None,
+            exclude_edges=exclude_edges,
             replace=replace,
             etype_offset=etype_offset,
             etype_sorted=etype_sorted,
@@ -1092,7 +1092,7 @@ def issue_remote_req(node_ids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
-            exclude_edges=None,
+            exclude_edges=exclude_edges,
             replace=replace,
             use_graphbolt=use_graphbolt,
         )
@@ -1108,7 +1108,7 @@ def local_access(local_g, partition_book, local_nids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
-            exclude_edges=None,
+            exclude_edges=exclude_edges,
             replace=replace,
         )
 

diff --git a/python/dgl/sampling/neighbor.py b/python/dgl/sampling/neighbor.py
@@ -159,6 +159,10 @@ def sample_etype_neighbors(
     As a result, users should avoid performing in-place operations
     on the node features of the new graph to avoid feature corruption.
     """
+    if exclude_edges is not None:
+        raise DGLError(
+            "exclude_edges is not supported for sample_etype_neighbors"
+        )
     if g.device != F.cpu():
         raise DGLError("The graph should be in cpu.")
     # (BarclayII) because the homogenized graph no longer contains the *name* of edge