Repository: dmlc/dgl

Issue:
GINConv Backward broken for a case in dgl 0.7.1+
## 🐛 Bug
GINConv backward is broken in the following script.

## To Reproduce

Steps to reproduce the behavior:

1. Run the code.

```python
import dgl.nn.pytorch as nn
import torch
import dgl


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin_1 = torch.nn.Linear(100, 100)
        self.gin = nn.GINConv(None, 'max')

    def forward(self, sg):
        x1 = torch.ones([32, 100])
        x2 = torch.ones([36, 100])
        x1 = self.lin_1(x1)  # have trainable weights so backward will compute
        g = self.gin(sg, (x1, x2), edge_weight=sg.edata['w'])
        return g.mean()


if __name__ == "__main__":
    net = Net()
    graph1 = dgl.rand_bipartite('a', 'con', 'b', 32, 36, 100)
    graph1.edges['con'].data['w'] = torch.ones([100, 1])
    pred = net(graph1)
    print(pred)
    pred.backward()

```

## Expected behavior

It should work.

## Environment

 - DGL Version (e.g., 1.0): 0.8a211021
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.9.1
 - OS (e.g., Linux): GNU/Linux 5.4.0-80-generic x86_64
 - How you installed DGL (`conda`, `pip`, source): `pip`
 - Build command you used (if compiling from source): N/A
 - Python version: 3.9.5
 - CUDA/cuDNN version (if applicable): 11.0
 - GPU models and configuration (e.g. V100): reproduces on CPU and V100
 - Any other relevant information: Does not reproduce with dgl 0.7.0 on Windows 11

## Additional context


GINConv Backward broken for a case in dgl 0.7.1+
## 🐛 Bug
GINConv backward is broken in the following script.

## To Reproduce

Steps to reproduce the behavior:

1. Run the code.

```python
import dgl.nn.pytorch as nn
import torch
import dgl


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin_1 = torch.nn.Linear(100, 100)
        self.gin = nn.GINConv(None, 'max')

    def forward(self, sg):
        x1 = torch.ones([32, 100])
        x2 = torch.ones([36, 100])
        x1 = self.lin_1(x1)  # have trainable weights so backward will compute
        g = self.gin(sg, (x1, x2), edge_weight=sg.edata['w'])
        return g.mean()


if __name__ == "__main__":
    net = Net()
    graph1 = dgl.rand_bipartite('a', 'con', 'b', 32, 36, 100)
    graph1.edges['con'].data['w'] = torch.ones([100, 1])
    pred = net(graph1)
    print(pred)
    pred.backward()

```

## Expected behavior

It should work.

## Environment

 - DGL Version (e.g., 1.0): 0.8a211021
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.9.1
 - OS (e.g., Linux): GNU/Linux 5.4.0-80-generic x86_64
 - How you installed DGL (`conda`, `pip`, source): `pip`
 - Build command you used (if compiling from source): N/A
 - Python version: 3.9.5
 - CUDA/cuDNN version (if applicable): 11.0
 - GPU models and configuration (e.g. V100): reproduces on CPU and V100
 - Any other relevant information: Does not reproduce with dgl 0.7.0 on Windows 11

## Additional context




Discussion:
Update: perhaps it is related with PyTorch version. It reproduces with dgl 0.7.1 on the linux server too.
Update: it reproduces with PyTorch 1.8.0 and dgl 0.7.1 or 0.8.0a211021. It does not reproduce with dgl 0.7.0. So it is a commit between 0.7.0 and 0.7.1 that breaks.
Update: perhaps it is related with PyTorch version. It reproduces with dgl 0.7.1 on the linux server too.
Update: it reproduces with PyTorch 1.8.0 and dgl 0.7.1 or 0.8.0a211021. It does not reproduce with dgl 0.7.0. So it is a commit between 0.7.0 and 0.7.1 that breaks.

Resolving patch:
diff --git a/python/dgl/backend/pytorch/sparse.py b/python/dgl/backend/pytorch/sparse.py
--- a/python/dgl/backend/pytorch/sparse.py
+++ b/python/dgl/backend/pytorch/sparse.py
@@ -102,7 +102,7 @@ def spmm_cache_Y(binary_op, reduce_op, req_grad_X, req_grad_Y):
 
 def spmm_cache_argX(binary_op, reduce_op, req_grad_X, req_grad_Y):
     """Rules to identify whether to cache argX in SpMM forward stage."""
-    if req_grad_X:
+    if req_grad_X or req_grad_Y:
         if reduce_op in ['min', 'max']:
             return True
     return False
@@ -110,7 +110,7 @@ def spmm_cache_argX(binary_op, reduce_op, req_grad_X, req_grad_Y):
 
 def spmm_cache_argY(binary_op, reduce_op, req_grad_X, req_grad_Y):
     """Rules to identify whether to cache argY in SpMM forward stage."""
-    if req_grad_Y:
+    if req_grad_X or req_grad_Y:
         if reduce_op in ['min', 'max']:
             return True
     return False


Test patch: