Repository: pyg-team/pytorch_geometric (python)

Issue:
TorchScript instance from torch.jit.script() runs successfully but throws type mismatch when loading from file
### 🐛 Describe the bug

In my example below, I can compile a Torchscript model of GENConv and compute the result Y1, however when I load the model from the saved file and try to calculate Y2 it throws a Type mismatch error. This is related to the input validation of Union[Tensor, OptPairTensor] in the Message Passing class. Torchscript requires the type to be the same for the true/false branches.

I'm not sure how to resolve this and am happy to contribute if someone can point me in the right direction.

I'm using the latest nightly build of pyg.

```
import torch
from torch import nn, Tensor
from torch_geometric.typing import Adj, OptTensor, Union, OptPairTensor
from torch_geometric.nn.conv import GENConv, GraphConv

class MyConv(nn.Module):
    def __init__(self):
        super(MyConv, self).__init__()
        self.conv = GENConv(64,
                            64,
                            learn_p=True,
                            num_layers=1)

    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_attr: OptTensor = None):
        y = self.conv(x, edge_index, edge_attr)
        return y

if __name__ == "__main__":
    device = 'cpu'
    save_path = "example.pt"
    model = MyConv()
    scripted_model = torch.jit.script(model)
    scripted_model.save(save_path)

    n_nodes = 200
    n_edges = 100
    x = torch.rand((n_nodes, 64), dtype=torch.float32, device=device)
    edge_index = torch.randint(2, (2, n_edges), dtype=torch.int64, device=device)
    edge_attr = torch.rand((n_edges, 64), dtype=torch.float32, device=device)
   
    y1 = scripted_model(x=x,
                                     edge_index=edge_index,
                                     edge_attr=edge_attr)

    model_from_file = torch.jit.load(save_path)
    y2 = model_from_file(x=x,
                                       edge_index=edge_index,
                                       edge_attr=edge_attr)
```

**Error:**
```
/miniforge3/envs/pyg/bin/python ...../Example.py 
Traceback (most recent call last):
  File "......../Example.py", line 38, in <module>
    model_from_file = torch.jit.load(save_path)
  File "......../lib/python3.10/site-packages/torch/jit/_serialization.py", line 162, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files, _restore_shapes)  # type: ignore[call-arg]
RuntimeError: 

Type mismatch: edge_index0 is set to type __torch__.torch_geometric.typing.SparseTensor (of Python compilation unit at: 0x561bb9716290) in the true branch and type Tensor in the false branch:
  File "code/__torch__/torch_geometric/nn/conv/gen_conv.py", line 52
      x0 = x2
    _7 = isinstance(edge_index, __torch__.torch_geometric.typing.SparseTensor)
    if _7:
    ~~~~~~
      edge_index1 = unchecked_cast(__torch__.torch_geometric.typing.SparseTensor, edge_index)
      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      edge_attr1 = ((edge_index1).__storage_getter()).value()
      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      edge_attr0, edge_index0 = edge_attr1, edge_index1
      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    else:
    ~~~~~
      edge_index2 = unchecked_cast(Tensor, edge_index)
      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      if _4(edge_index2, ):
      ~~~~~~~~~~~~~~~~~~~~~
        _8, value, = _5(edge_index2, )
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        if torch.gt(torch.dim(value), 1):
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
          _9 = True
          ~~~~~~~~~
        else:
        ~~~~~
          _10 = torch.__not__(bool(torch.all(value)))
          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
          _9 = _10
          ~~~~~~~~
        if _9:
        ~~~~~~
          edge_attr3 : Optional[Tensor] = value
          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        else:
        ~~~~~
          edge_attr3 = edge_attr
          ~~~~~~~~~~~~~~~~~~~~~~
        edge_attr2 : Optional[Tensor] = edge_attr3
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      else:
      ~~~~~
        edge_attr2 = edge_attr
        ~~~~~~~~~~~~~~~~~~~~~~
      edge_attr0, edge_index0 = edge_attr2, edge_index2
      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    if torch.__isnot__(edge_attr0, None):
      edge_attr5 = unchecked_cast(Tensor, edge_attr0)
and was used here:
  File "code/__torch__/torch_geometric/nn/conv/gen_conv.py", line 88
    else:
      edge_attr6 = edge_attr4
    out = (self).propagate(edge_index0, x0, edge_attr6, size, )
                           ~~~~~~~~~~~ <--- HERE
    x_dst = (x0)[1]
    if torch.__isnot__(x_dst, None):


Process finished with exit code 1
```



### Versions

pyg-nightly         2.4.0.dev20240205

pytorch                   2.1.2           py3.10_cuda12.1_cudnn8.9.2_0    pytorch
pytorch-cluster           1.6.3           py310_torch_2.1.0_cu121    pyg
pytorch-cuda              12.1                 ha16c6d3_5    pytorch
pytorch-lightning         2.1.3              pyhd8ed1ab_0    conda-forge
pytorch-mutex             1.0                        cuda    pytorch
pytorch-scatter           2.1.2           py310_torch_2.1.0_cu121    pyg
onnx                      1.15.0          py310heee6c8b_1    conda-forge
onnxruntime               1.16.3          py310hd4b7fbc_4_cpu    conda-forge
cuda-opencl               12.3.101                      0    nvidia
cuda-runtime              12.1.0                        0    nvidia

Pull request description:
Fix `torch.jit.load` for `MessagePassing` modules
Fixes #8867

Patch:
diff --git a/torch_geometric/nn/conv/collect.jinja b/torch_geometric/nn/conv/collect.jinja
index e2f4c1ee1..7c1ad2fb3 100644
--- a/torch_geometric/nn/conv/collect.jinja
+++ b/torch_geometric/nn/conv/collect.jinja
@@ -30,34 +30,43 @@ def {{collect_name}}(
     i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
 
     # Collect special arguments:
-    if isinstance(edge_index, Tensor) and is_torch_sparse_tensor(edge_index):
+    if isinstance(edge_index, Tensor):
+        if is_torch_sparse_tensor(edge_index):
 {%- if 'edge_index' in collect_param_dict %}
-        raise ValueError("Cannot collect 'edge_indices' for sparse matrices")
+            raise ValueError("Cannot collect 'edge_indices' for sparse matrices")
 {%- endif %}
-        adj_t = edge_index
-        if adj_t.layout == torch.sparse_coo:
-            edge_index_i = adj_t.indices()[0]
-            edge_index_j = adj_t.indices()[1]
-            ptr = None
-        elif adj_t.layout == torch.sparse_csr:
-            ptr = adj_t.crow_indices()
-            edge_index_j = adj_t.col_indices()
-            edge_index_i = ptr2index(ptr, output_size=edge_index_j.numel())
-        else:
-            raise ValueError(f"Received invalid layout '{adj_t.layout}'")
+            adj_t = edge_index
+            if adj_t.layout == torch.sparse_coo:
+                edge_index_i = adj_t.indices()[0]
+                edge_index_j = adj_t.indices()[1]
+                ptr = None
+            elif adj_t.layout == torch.sparse_csr:
+                ptr = adj_t.crow_indices()
+                edge_index_j = adj_t.col_indices()
+                edge_index_i = ptr2index(ptr, output_size=edge_index_j.numel())
+            else:
+                raise ValueError(f"Received invalid layout '{adj_t.layout}'")
 
 {%- if 'edge_weight' in collect_param_dict %}
-        if edge_weight is None:
-            edge_weight = adj_t.values()
+            if edge_weight is None:
+                edge_weight = adj_t.values()
 {%- elif 'edge_attr' in collect_param_dict %}
-        if edge_attr is None:
-            _value = adj_t.values()
-            edge_attr = None if _value.dim() == 1 else _value
+            if edge_attr is None:
+                _value = adj_t.values()
+                edge_attr = None if _value.dim() == 1 else _value
 {%- elif 'edge_type' in collect_param_dict %}
-        if edge_type is None:
-            edge_type = adj_t.values()
+            if edge_type is None:
+                edge_type = adj_t.values()
 {%- endif %}
 
+        else:
+{%- if 'adj_t' in collect_param_dict %}
+            raise ValueError("Cannot collect 'adj_t' for edge indices")
+{%- endif %}
+            edge_index_i = edge_index[i]
+            edge_index_j = edge_index[j]
+            ptr = None
+
     elif isinstance(edge_index, SparseTensor):
 {%- if 'edge_index' in collect_param_dict %}
         raise ValueError("Cannot collect 'edge_indices' for sparse matrices")
@@ -77,13 +86,8 @@ def {{collect_name}}(
             edge_type = _value
 {%- endif %}
 
-    elif isinstance(edge_index, Tensor):
-{%- if 'adj_t' in collect_param_dict %}
-        raise ValueError("Cannot collect 'adj_t' for edge indices")
-{%- endif %}
-        edge_index_i = edge_index[i]
-        edge_index_j = edge_index[j]
-        ptr = None
+    else:
+        raise NotImplementedError
 
 {%- if 'edge_weight' in collect_param_dict and
     collect_param_dict['edge_weight'].type_repr.endswith('Tensor') %}
diff --git a/torch_geometric/nn/conv/gen_conv.py b/torch_geometric/nn/conv/gen_conv.py
index 0b297a77f..4127088a0 100644
--- a/torch_geometric/nn/conv/gen_conv.py
+++ b/torch_geometric/nn/conv/gen_conv.py
@@ -15,14 +15,7 @@ from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.nn.inits import reset
 from torch_geometric.nn.norm import MessageNorm
-from torch_geometric.typing import (
-    Adj,
-    OptPairTensor,
-    OptTensor,
-    Size,
-    SparseTensor,
-)
-from torch_geometric.utils import is_torch_sparse_tensor, to_edge_index
+from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size
 
 
 class MLP(Sequential):
@@ -216,20 +209,6 @@ class GENConv(MessagePassing):
         if hasattr(self, 'lin_src'):
             x = (self.lin_src(x[0]), x[1])
 
-        if isinstance(edge_index, SparseTensor):
-            edge_attr = edge_index.storage.value()
-        elif is_torch_sparse_tensor(edge_index):
-            _, value = to_edge_index(edge_index)
-            if value.dim() > 1 or not value.all():
-                edge_attr = value
-
-        if edge_attr is not None and hasattr(self, 'lin_edge'):
-            edge_attr = self.lin_edge(edge_attr)
-
-        # Node and edge feature dimensionalites need to match.
-        if edge_attr is not None:
-            assert x[0].size(-1) == edge_attr.size(-1)
-
         # propagate_type: (x: OptPairTensor, edge_attr: OptTensor)
         out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)
 
@@ -250,6 +229,12 @@ class GENConv(MessagePassing):
         return self.mlp(out)
 
     def message(self, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
+        if edge_attr is not None and hasattr(self, 'lin_edge'):
+            edge_attr = self.lin_edge(edge_attr)
+
+        if edge_attr is not None:
+            assert x_j.size(-1) == edge_attr.size(-1)
+
         msg = x_j if edge_attr is None else x_j + edge_attr
         return msg.relu() + self.eps
 


Test patch:
diff --git a/test/nn/conv/test_message_passing.py b/test/nn/conv/test_message_passing.py
index 65e99eb58..2bd285282 100644
--- a/test/nn/conv/test_message_passing.py
+++ b/test/nn/conv/test_message_passing.py
@@ -1,4 +1,5 @@
 import copy
+import os.path as osp
 from typing import Tuple, Union
 
 import pytest
@@ -184,6 +185,15 @@ def test_my_conv_jit():
         jit.fuse = True
 
 
+def test_my_conv_jit_save(tmp_path):
+    path = osp.join(tmp_path, 'model.pt')
+
+    conv = MyConv(8, 32)
+    conv = torch.jit.script(conv)
+    torch.jit.save(conv, path)
+    conv = torch.jit.load(path)
+
+
 @pytest.mark.parametrize('aggr', ['add', 'sum', 'mean', 'min', 'max', 'mul'])
 def test_my_conv_aggr(aggr):
     x = torch.randn(4, 8)