fix and lots of tests (dmlc#2650)

CyrilZhao-sudo · Feb 12, 2021 · 9e63010 · 9e63010
1 parent cf8a3fb
commit 9e63010
Show file tree

Hide file tree

Showing 6 changed files with 322 additions and 262 deletions.
diff --git a/python/dgl/nn/mxnet/utils.py b/python/dgl/nn/mxnet/utils.py
@@ -84,7 +84,7 @@ def bmm_maybe_select(A, B, index):
         return B[index, A, :]
     else:
         BB = nd.take(B, index, axis=0)
-        return nd.batch_dot(A.expand_dims(1), BB).squeeze()
+        return nd.batch_dot(A.expand_dims(1), BB).squeeze(1)
 
 def normalize(x, p=2, axis=1, eps=1e-12):
     r"""Performs :math:`L_p` normalization of inputs over specified dimension.

diff --git a/python/dgl/nn/pytorch/conv/relgraphconv.py b/python/dgl/nn/pytorch/conv/relgraphconv.py
@@ -238,7 +238,7 @@ def basis_message_func(self, edges, etypes):
                 etypes = th.repeat_interleave(th.arange(len(etypes), device=device),
                                               th.tensor(etypes, device=device))
             weight = weight.index_select(0, etypes)
-            msg = th.bmm(h.unsqueeze(1), weight).squeeze()
+            msg = th.bmm(h.unsqueeze(1), weight).squeeze(1)
 
         if 'norm' in edges.data:
             msg = msg * edges.data['norm']

diff --git a/python/dgl/nn/tensorflow/utils.py b/python/dgl/nn/tensorflow/utils.py
@@ -87,7 +87,7 @@ def bmm_maybe_select(A, B, index):
         return tf.gather(B, flatidx)
     else:
         BB = tf.gather(B, index)
-        return tf.squeeze(tf.matmul(tf.expand_dims(A, 1), BB))
+        return tf.squeeze(tf.matmul(tf.expand_dims(A, 1), BB), 1)
 
 
 class Identity(layers.Layer):