test dist_server

alibaba · Oct 28, 2024 · df28ab4 · df28ab4
1 parent 0d7bf9e
commit df28ab4
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 202 deletions.
diff --git a/graphlearn_torch/python/distributed/dist_server.py b/graphlearn_torch/python/distributed/dist_server.py
@@ -109,22 +109,10 @@ def get_edge_index(self, edge_type, layout):
     if layout == 'coo':
       row, col, _, _ = graph.topo.to_coo()
       result = (row, col)
-      size = (graph.row_count, graph.col_count)
-    elif layout == 'csr':
-      row, col, _, _ = graph.topo.to_csr()
-      result = (row, col)
-      size = (graph.row_count, col.max().item()+1)
-    elif layout == 'csc':
-      row, col, _, _ = graph.topo.to_csc()
-      result = (row, col)
-      size = (row.max().item()+1, graph.row_count)
     else:
       raise ValueError(f"Invalid layout {layout}")
-    return result, size
-
-  def get_edge_size(self, edge_type, layout):
-    _, size = self.get_edge_index(edge_type, layout)
-    return size
+    return result
+
 
   def create_sampling_producer(
     self,