Repository: pyg-team/pytorch_geometric
Pull request: Support `SparseTensor` label in `LightGCN`
fix the bug of models/lightgcn when edge_indx is an instance of SparseTensor in forward method

Problem statement:
models/lightgcn raises an error when the `edge_index` is an instance of SparseTensor in forward method
### 🐛 Describe the bug

 When `edge_index ` is an instance of SparseTensor, method `forward` as follows would raise an error:
 ```
  def forward(self, edge_index: Adj,
                edge_label_index: OptTensor = None) -> Tensor:

        if edge_label_index is None:

                edge_label_index = edge_index

        out = self.get_embedding(edge_index)

        out_src = out[edge_label_index[0]]
        out_dst = out[edge_label_index[1]]
        return (out_src * out_dst).sum(dim=-1)

```


### Environment

* PyG version:2.0.2
* PyTorch version: 1.7.1
* OS: CentOS
* Python version: 3.8
* CUDA/cuDNN version:11.0
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip 
* Any other relevant information (*e.g.*, version of `torch-scatter`):

Hints from issue discussion:

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/torch_geometric/nn/models/lightgcn.py b/torch_geometric/nn/models/lightgcn.py
index 5cbda80fbec8..f497cad0dd53 100644
--- a/torch_geometric/nn/models/lightgcn.py
+++ b/torch_geometric/nn/models/lightgcn.py
@@ -5,6 +5,7 @@
 from torch import Tensor
 from torch.nn import Embedding, ModuleList
 from torch.nn.modules.loss import _Loss
+from torch_sparse import SparseTensor
 
 from torch_geometric.nn.conv import LGConv
 from torch_geometric.typing import Adj, OptTensor
@@ -110,7 +111,21 @@ def forward(self, edge_index: Adj,
                 :obj:`edge_index` will be used instead. (default: :obj:`None`)
         """
         if edge_label_index is None:
-            edge_label_index = edge_index
+            if isinstance(edge_index, SparseTensor):
+                edge_label_index = edge_index.to_torch_sparse_coo_tensor(
+                ).coalesce().indices()
+            elif isinstance(edge_index, torch.Tensor) and edge_index.is_sparse(
+            ) and not edge_index.is_coalesced():
+                edge_label_index = edge_index.coalesce().indices()
+            elif isinstance(edge_index, torch.Tensor) and edge_index.is_sparse(
+            ) and edge_index.is_coalesced():
+                edge_label_index = edge_index.indices()
+            elif isinstance(edge_index,
+                            torch.Tensor) and not edge_index.is_sparse():
+                assert edge_index.shape == 2
+                edge_label_index = edge_index
+            else:
+                raise TypeError
 
         out = self.get_embedding(edge_index)
 


Code review comments:
diff_hunk: @@ -110,7 +111,21 @@ def forward(self, edge_index: Adj,
                 :obj:`edge_index` will be used instead. (default: :obj:`None`)
         """
         if edge_label_index is None:
-            edge_label_index = edge_index
+            if isinstance(edge_index, SparseTensor):
+                edge_label_index = edge_index.to_torch_sparse_coo_tensor(
+                ).coalesce().indices()
+            elif isinstance(edge_index, torch.Tensor) and edge_index.is_sparse(
path: torch_geometric/nn/models/lightgcn.py
text: @user1:
I don‘t think we need to support this case right now. Any `tensor.is_sparse` logic should be dropped here.

diff_hunk: @@ -110,7 +111,21 @@ def forward(self, edge_index: Adj,
                 :obj:`edge_index` will be used instead. (default: :obj:`None`)
         """
         if edge_label_index is None:
-            edge_label_index = edge_index
+            if isinstance(edge_index, SparseTensor):
+                edge_label_index = edge_index.to_torch_sparse_coo_tensor(
path: torch_geometric/nn/models/lightgcn.py
text: @user1:
```suggestion
                edge_label_index = torch.stack(edge_index.coo()[:2], dim=0)
```