[BUGFIX] fix a bug in graph partitioning. (#1769)

* fix

* use utils.toindex in the right place.

* fix.

* update tensor for mxnet backend.

* fix

* fix

python/dgl/backend/mxnet/tensor.py

@@ -35,14 +35,23 @@ def cpu():
     return mx.cpu()
 
 def tensor(data, dtype=None):
-    # MXNet always returns a float tensor regardless of type inside data.
-    # This is a workaround.
-    if dtype is None:
-        if isinstance(data[0], numbers.Integral):
-            dtype = np.int64
+    if isinstance(data, nd.NDArray):
+        if dtype is None or data.dtype == dtype:
+            return data
         else:
-            dtype = np.float32
-    return nd.array(data, dtype=dtype)
+            return nd.cast(data, dtype)
+    else:
+        if dtype is None:
+            if isinstance(data, numbers.Number):
+                dtype = np.int64 if isinstance(data, numbers.Integral) else np.float32
+            elif isinstance(data, np.ndarray):
+                dtype = data.dtype
+                # mxnet doesn't support bool
+                if dtype == np.bool:
+                    dtype = np.int32
+            else:
+                dtype = np.int64 if isinstance(data[0], numbers.Integral) else np.float32
+        return nd.array(data, dtype=dtype)
 
 def as_scalar(data):
     return data.asscalar()

python/dgl/backend/pytorch/tensor.py

@@ -31,7 +31,7 @@ def cpu():
     return th.device('cpu')
 
 def tensor(data, dtype=None):
-    return th.tensor(data, dtype=dtype)
+    return th.as_tensor(data, dtype=dtype)
 
 def as_scalar(data):
     return data.item()
@@ -272,6 +272,8 @@ def clone(input):
     return input.clone()
 
 def unique(input):
+    if input.dtype == th.bool:
+        input = input.type(th.int8)
     return th.unique(input)
 
 def full_1d(length, fill_value, dtype, ctx):

python/dgl/distributed/dist_graph.py

@@ -576,8 +576,7 @@ def _get_overlap(mask_arr, ids):
         masks = mask_arr[ids]
         return F.boolean_mask(ids, masks)
     else:
-        mask_arr = utils.toindex(mask_arr)
-        masks = F.gather_row(mask_arr.tousertensor(), ids)
+        masks = F.gather_row(F.tensor(mask_arr), ids)
         return F.boolean_mask(ids, masks)
 
 def _split_local(partition_book, rank, elements, local_eles):
@@ -615,8 +614,7 @@ def _split_even(partition_book, rank, elements):
         # I hope it's OK.
         eles = F.nonzero_1d(elements[0:len(elements)])
     else:
-        elements = utils.toindex(elements)
-        eles = F.nonzero_1d(elements.tousertensor())
+        eles = F.nonzero_1d(F.tensor(elements))
 
     # here we divide the element list as evenly as possible. If we use range partitioning,
     # the split results also respect the data locality. Range partitioning is the default

python/dgl/distributed/graph_partition_book.py

@@ -100,12 +100,10 @@ class GraphPartitionBook:
         assert num_parts > 0, 'num_parts must be greater than zero.'
         self._part_id = int(part_id)
         self._num_partitions = int(num_parts)
-        node_map = utils.toindex(node_map)
-        self._nid2partid = node_map.tousertensor()
+        self._nid2partid = F.tensor(node_map)
         assert F.dtype(self._nid2partid) in (F.int32, F.int64), \
                 'the node map must be stored in an integer array'
-        edge_map = utils.toindex(edge_map)
-        self._eid2partid = edge_map.tousertensor()
+        self._eid2partid = F.tensor(edge_map)
         assert F.dtype(self._eid2partid) in (F.int32, F.int64), \
                 'the edge map must be stored in an integer array'
         # Get meta data of the partition book.
@@ -374,10 +372,12 @@ class RangePartitionBook:
         assert num_parts > 0, 'num_parts must be greater than zero.'
         self._partid = part_id
         self._num_partitions = num_parts
-        node_map = utils.toindex(node_map)
-        edge_map = utils.toindex(edge_map)
-        self._node_map = node_map.tonumpy()
-        self._edge_map = edge_map.tonumpy()
+        if not isinstance(node_map, np.ndarray):
+            node_map = F.asnumpy(node_map)
+        if not isinstance(edge_map, np.ndarray):
+            edge_map = F.asnumpy(edge_map)
+        self._node_map = node_map
+        self._edge_map = edge_map
         # Get meta data of the partition book
         self._partition_meta_data = []
         for partid in range(self._num_partitions):

python/dgl/transform.py

@@ -701,8 +701,7 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False):
     if balance_ntypes is not None:
         assert len(balance_ntypes) == g.number_of_nodes(), \
                 "The length of balance_ntypes should be equal to #nodes in the graph"
-        balance_ntypes = utils.toindex(balance_ntypes)
-        balance_ntypes = balance_ntypes.tousertensor()
+        balance_ntypes = F.tensor(balance_ntypes)
         uniq_ntypes = F.unique(balance_ntypes)
         for ntype in uniq_ntypes:
             vwgt.append(F.astype(balance_ntypes == ntype, F.int64))