Repository: dmlc/dgl

Issue:
[Bug] inconsistent behavior between update_all and multi_update_all for heterogeneous graphs with only one ntype
## 🐛 Bug

Similar to #4816, update_all is not working for heterogeneous graphs with only one ntype, causing an error in backward. 
However, using multi_update_all is fine.

## To Reproduce

Steps to reproduce the behavior:

```python
import dgl
import torch as th
import dgl.function as fn

graph_data = {
   ('drug', 'interacts', 'drug'): (th.tensor([0, 1]), th.tensor([1, 2])),
   ('drug', 'treats', 'drug'): (th.tensor([1]), th.tensor([2]))
}

g = dgl.heterograph(graph_data)

g.srcdata['h'] = th.nn.Parameter(th.zeros(3)+5.)
g.edges['interacts'].data['w'] = th.ones(2)
g.edges['treats'].data['w'] = th.ones(1)
for etype in g.canonical_etypes:
    g.apply_edges(fn.u_mul_e('h', 'w', 'm'), etype=etype)

g.update_all(fn.copy_e('m','t'), fn.sum('t', 'h1'))
#g.multi_update_all({'interacts': (fn.copy_e('m','t'), fn.sum('t', 'h1')), 'treats': (fn.copy_e('m', 't'), fn.sum('t', 'h1'))}, "sum") # this line works

print(g.ndata['h1'].sum())
print(g.ndata['h1'].sum().backward())

```

## Expected behavior

```
    print(g.ndata['h1'].sum().backward())
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/torch/_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/torch/autograd/function.py", line 199, in apply
    return user_fn(self, *args)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 111, in decorate_bwd
    return bwd(*args, **kwargs)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py", line 253, in backward
    dY = gsddmm_hetero(gidx, 'copy_rhs', X_len, 'u', 'v', *tpl_X_dZ)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py", line 761, in gsddmm_hetero
    return GSDDMM_hetero.apply(g, op, lhs_len, lhs_target, rhs_target, *lhs_and_rhs_tuple)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py", line 354, in forward
    out = _gsddmm_hetero(gidx, op, X_len, lhs_target, rhs_target, feats)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/dgl/sparse.py", line 542, in _gsddmm_hetero
    rhs = F.unsqueeze(lhs, -1)
  File "/home/ec2-user/anaconda3/envs/py39/lib/python3.9/site-packages/dgl/backend/pytorch/tensor.py", line 215, in unsqueeze
    return th.unsqueeze(input, dim)
TypeError: unsqueeze(): argument 'input' (position 1) must be Tensor, not NoneType
```

## Environment

 - DGL Version (e.g., 1.0): 0.9.1
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.10.0
 - OS (e.g., Linux): Linux
 - How you installed DGL (`conda`, `pip`, source): pip
 - Python version:  3.9.7
 - CUDA/cuDNN version (if applicable): 102
 - GPU models and configuration (e.g. V100): T4



Discussion:

Resolving patch:
diff --git a/python/dgl/sparse.py b/python/dgl/sparse.py
--- a/python/dgl/sparse.py
+++ b/python/dgl/sparse.py
@@ -605,7 +605,7 @@ def _gsddmm_hetero(
                 expand_lhs = True
         if use_rhs:
             if rhs is not None and F.ndim(rhs) == 1:
-                rhs = F.unsqueeze(lhs, -1)
+                rhs = F.unsqueeze(rhs, -1)
                 expand_rhs = True
         ctx = F.context(lhs) if use_lhs else F.context(rhs)
         dtype = F.dtype(lhs) if use_lhs else F.dtype(rhs)


Test patch: