fix edge_softmax (#2160)

Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>

python/dgl/backend/mxnet/sparse.py

@@ -269,7 +269,7 @@ class EdgeSoftmax(mx.autograd.Function):
     def __init__(self, gidx, eids, norm_by):
         super(EdgeSoftmax, self).__init__()
         if not is_all(eids):
-            gidx = gidx.edge_subgraph(eids.astype(gidx.dtype), True)
+            gidx = gidx.edge_subgraph([eids], True).graph
         if norm_by == 'src':
             gidx = gidx.reverse()
         self.gidx = gidx

python/dgl/backend/pytorch/sparse.py

@@ -196,7 +196,7 @@ class EdgeSoftmax(th.autograd.Function):
         # remember to save the graph to backward cache before making it
         # a local variable
         if not is_all(eids):
-            gidx = gidx.edge_subgraph(eids.type(gidx.dtype), True)
+            gidx = gidx.edge_subgraph([eids], True).graph
         if norm_by == 'src':
             gidx = gidx.reverse()
         score_max = _gspmm(gidx, 'copy_rhs', 'max', None, score)[0]

python/dgl/backend/tensorflow/sparse.py

@@ -225,7 +225,7 @@ def gsddmm(gidx, op, X, Y, lhs_target='u', rhs_target='v'):
 
 def edge_softmax_real(gidx, score, eids=ALL, norm_by='dst'):
     if not is_all(eids):
-        gidx = gidx.edge_subgraph(tf.cast(eids, gidx.dtype), True)
+        gidx = gidx.edge_subgraph([eids], True).graph
     if norm_by == 'src':
         gidx = gidx.reverse()
     score_max = _gspmm(gidx, 'copy_rhs', 'max', None, score)[0]

python/dgl/ops/edge_softmax.py

 """dgl edge_softmax operator module."""
 from ..backend import edge_softmax as edge_softmax_internal
-from ..base import ALL
+from ..backend import astype
+from ..base import ALL, is_all
 
 __all__ = ['edge_softmax']
 
@@ -103,5 +104,7 @@ def edge_softmax(graph, logits, eids=ALL, norm_by='dst'):
                 [1.0000],
                 [0.5000]])
     """
+    if not is_all(eids):
+        eids = astype(eids, graph.idtype)
     return edge_softmax_internal(graph._graph, logits,
                                  eids=eids, norm_by=norm_by)