Repository: dmlc/dgl

Issue:
AssertionError when tape.gradient (Tensorflow)
## ❓ Questions and Help

Hi, I'm trying to train a nn model with tf backend but I get AssertionError during tape.gradient. What I got stuck can be reproduced as following simple script:
```python
graph = dgl.graph(([0,1],[2,3]))
v = tf.Variable(0.1)
with tf.GradientTape() as tape:
    graph.edata["w"] = tf.repeat(v, graph.number_of_edges())
    graph.update_all(fn.copy_e("w", "m"), fn.max("m", "ws"))
    pred = graph.ndata["ws"]
    loss = pred - .1 # just for test, should work tho
tape.gradient(loss, [v])
```

Output:

> AssertionError                            Traceback (most recent call last)
 in 
      6     pred = graph.ndata["ws"]
      7     loss = pred - .1
----> 8 tape.gradient(loss, [v])
~/.conda/envs/.../lib/python3.7/site-packages/tensorflow/python/eager/backprop.py in gradient(self, target, sources, output_gradients, unconnected_gradients)
   1071         output_gradients=output_gradients,
   1072         sources_raw=flat_sources_raw,
-> 1073         unconnected_gradients=unconnected_gradients)
   1074 
   1075     if not self._persistent:
~/.conda/envs/.../lib/python3.7/site-packages/tensorflow/python/eager/imperative_grad.py in imperative_grad(tape, target, sources, output_gradients, sources_raw, unconnected_gradients)
     75       output_gradients,
     76       sources_raw,
---> 77       compat.as_str(unconnected_gradients.value))
~/.conda/envs/.../lib/python3.7/site-packages/tensorflow/python/ops/custom_gradient.py in actual_grad_fn(*result_grads)
    444                          "@custom_gradient grad_fn.")
    445     else:
--> 446       input_grads = grad_fn(*result_grads)
    447       variable_grads = []
    448     flat_grads = nest.flatten(input_grads)
~/.conda/envs/.../lib/python3.7/site-packages/dgl/backend/tensorflow/sparse.py in grad(dZ)
    142                     if op == 'div': dY = -dY / (Y ** 2)
    143                 elif op in ['add', 'sub', 'copy_rhs']:
--> 144                     dY = _scatter_nd(argY, _addsub(op, dZ), Y.shape[0])
    145             dY = _reduce_grad(dY, Y.shape)
    146         else:
~/.conda/envs/.../lib/python3.7/site-packages/dgl/backend/tensorflow/sparse.py in _scatter_nd(index, src, n_rows)
      9 
     10 def _scatter_nd(index, src, n_rows):
---> 11     assert index.shape == src.shape
     12     shp = index.shape
     13     ctx = context(src)
AssertionError: 

It might be a bug but I'm not sure I'm doing right with the API. If I replace `fn.max` to `fn.mean` then it works fine and I expect the same result with `fn.max` and `fn.min` but they seem not to work. I'm using 0.5.2. Thanks for help in advance. 


Discussion:
Thanks for your report. Confirmed it's a bug in our computation kernel on tensorflow.

Currently you can work it around by expand_dims to the one-dim tensor, such as changing to
```python
graph.edata["w"] = tf.expand_dims(tf.repeat(v, graph.number_of_edges()), -1)
```
in your posted code

Resolving patch:
diff --git a/python/dgl/backend/mxnet/sparse.py b/python/dgl/backend/mxnet/sparse.py
--- a/python/dgl/backend/mxnet/sparse.py
+++ b/python/dgl/backend/mxnet/sparse.py
@@ -25,7 +25,10 @@ def _scatter_nd(index, src, n_rows):
         offsets.append(
             (stride * offset_i).reshape((1,) * i + (di,) + (1,) * (ndim - 1 - i)))
         stride *= di
-    new_idx = index * stride + sum(offsets)
+    if ndim > 1:
+        new_idx = index * stride + sum(offsets)
+    else:
+        new_idx = index 
     src = src.reshape(-1)
     new_idx = new_idx.reshape(-1)
     rst = np.zeros((stride * n_rows,), dtype=src.dtype)
@@ -48,7 +51,10 @@ def _gather_nd(index, src):
         offsets.append(
             (stride * offset_i).reshape((1,) * i + (di,) + (1,) * (ndim - 1 - i)))
         stride *= di
-    new_idx = index * stride + copy_to(sum(offsets), ctx)
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
     src = src.reshape(-1)
     new_idx = new_idx.reshape(-1)
     rst = nd.take(src, new_idx).reshape(shp)
diff --git a/python/dgl/backend/tensorflow/sparse.py b/python/dgl/backend/tensorflow/sparse.py
--- a/python/dgl/backend/tensorflow/sparse.py
+++ b/python/dgl/backend/tensorflow/sparse.py
@@ -20,7 +20,10 @@ def _scatter_nd(index, src, n_rows):
         offsets.append(
             tf.reshape((stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)))
         stride *= di
-    new_idx = index * stride + copy_to(sum(offsets), ctx)
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
     src = tf.reshape(src, (-1,))
     new_idx = tf.reshape(new_idx, (-1, 1))
     rst = tf.reshape(tf.scatter_nd(new_idx, src, (stride * n_rows,)), (n_rows, *shp[1:]))
@@ -39,7 +42,10 @@ def _gather_nd(index, src):
         offsets.append(
             tf.reshape((stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)))
         stride *= di
-    new_idx = index * stride + copy_to(sum(offsets), ctx)
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
     src = tf.reshape(src, (-1,))
     new_idx = tf.reshape(new_idx, (-1))
     rst = tf.reshape(tf.gather(src, new_idx), shp)
@@ -72,7 +78,7 @@ def _reduce_grad(grad, shape):
     reduce_idx = np.asarray(np.nonzero(np.asarray(grad_shape) - np.asarray(in_shape)))
     reduce_idx += 1  # skip batch dim
     reduce_idx_tensor = tf.constant(tuple(
-        reduce_idx.flatten().tolist()))
+        reduce_idx.flatten().tolist()), dtype=tf.int32)
     grad = tf.reduce_sum(grad, axis=reduce_idx_tensor, keepdims=True)
     return tf.reshape(grad, shape)
 
diff --git a/python/dgl/sparse.py b/python/dgl/sparse.py
--- a/python/dgl/sparse.py
+++ b/python/dgl/sparse.py
@@ -167,6 +167,10 @@ def _gspmm(gidx, op, reduce_op, u, e):
     # To deal with scalar node/edge features.
     if (expand_u or not use_u) and (expand_e or not use_e):
         v = F.squeeze(v, -1)
+    if expand_u and use_cmp:
+        arg_u = F.squeeze(arg_u, -1)
+    if expand_e and use_cmp:
+        arg_e = F.squeeze(arg_e, -1)
     return v, (arg_u, arg_e)
 
 


Test patch:
diff --git a/tests/compute/test_sparse.py b/tests/compute/test_sparse.py
--- a/tests/compute/test_sparse.py
+++ b/tests/compute/test_sparse.py
@@ -79,7 +79,8 @@ def foo(edges):
     ((3, 3), (1, 3)),
     ((1,), (3,)),
     ((3,), (1,)),
-    ((1,), (1,))
+    ((1,), (1,)),
+    ((), ())
 ]
 
 sddmm_shapes = [