Repository: pyg-team/pytorch_geometric
Pull request: `JumpingKnowledge` Documentation Changes
Reference to issue #5499. Small improvement in the Jumping Knowledge code :)

Problem statement:
Small Documentation Change in Jumping Knowledge
### 📚 Describe the documentation issue

Sometimes when you have really complex networks, you need information about the architecture to be available quickly. Jumping knowledge difficults this when you select the lstm option, as from a simple print line you are not able to access the channels and layers selected. This small change should help people in tracking parameters in highly complex architectures.

### Suggest a potential alternative/fix

```
def __init__(self, mode: str, channels: Optional[int] = None,
                 num_layers: Optional[int] = None):
        super().__init__()
        self.mode = mode.lower()
        assert self.mode in ['cat', 'max', 'lstm']

        if mode == 'lstm':
            assert channels is not None, 'channels cannot be None for lstm'
            assert num_layers is not None, 'num_layers cannot be None for lstm'
            self.lstm = LSTM(channels, (num_layers * channels) // 2,
                             bidirectional=True, batch_first=True)
            self.att = Linear(2 * ((num_layers * channels) // 2), 1)
        else:
            self.lstm = None
            self.att = None

       self.reset_parameters()

def __repr__(self) -> str:
        return f'{self.__class__.__name__}({self.mode})'

```
This would be changed to
```
def __init__(self, mode: str, channels: Optional[int] = None,
                 num_layers: Optional[int] = None):
        super().__init__()
        self.mode = mode.lower()
        assert self.mode in ['cat', 'max', 'lstm']

        if mode == 'lstm':
            assert channels is not None, 'channels cannot be None for lstm'
            assert num_layers is not None, 'num_layers cannot be None for lstm'
            self.lstm = LSTM(channels, (num_layers * channels) // 2,
                             bidirectional=True, batch_first=True)
            self.att = Linear(2 * ((num_layers * channels) // 2), 1)
            self.channels = channels
            self.num_layers = num_layers
        else:
            self.lstm = None
            self.att = None
            self.channels = None
            self.num_layers = None

       self.reset_parameters()

def __repr__(self) -> str:
        if mode == 'lstm':
                return f'{self.__class__.__name__}({self.mode}, Channels: {self.channels}, Layers: {self.num_layers})'
        return f'{self.__class__.__name__}({self.mode})'

```

Hints from issue discussion:
Yes, this sounds good. Do you want to send a pull request for this? :)
Sure, It would be my first PR in a big repo tho, so bear with me if I make a mistake :) 
Don't worry :) Looking forward to it!

Commit message: JumpingKnowledge Documentation Changes

Patch to review:
diff --git a/test/nn/models/test_jumping_knowledge.py b/test/nn/models/test_jumping_knowledge.py
index f1f213a759e8..afa1008a47e3 100644
--- a/test/nn/models/test_jumping_knowledge.py
+++ b/test/nn/models/test_jumping_knowledge.py
@@ -29,7 +29,8 @@ def test_jumping_knowledge():
         assert torch.allclose(jit(xs), out)
 
     model = JumpingKnowledge('lstm', channels, num_layers)
-    assert model.__repr__() == 'JumpingKnowledge(lstm)'
+    assert model.__repr__() == 'JumpingKnowledge(lstm, Channels: ' +\
+                               f'{channels}, Layers: {num_layers})'
 
     out = model(xs)
     assert out.size() == (num_nodes, channels)
diff --git a/torch_geometric/nn/models/jumping_knowledge.py b/torch_geometric/nn/models/jumping_knowledge.py
index 1fddd9c2f1dd..c07a1f4aa268 100644
--- a/torch_geometric/nn/models/jumping_knowledge.py
+++ b/torch_geometric/nn/models/jumping_knowledge.py
@@ -51,9 +51,13 @@ def __init__(self, mode: str, channels: Optional[int] = None,
             self.lstm = LSTM(channels, (num_layers * channels) // 2,
                              bidirectional=True, batch_first=True)
             self.att = Linear(2 * ((num_layers * channels) // 2), 1)
+            self.channels = channels
+            self.num_layers = num_layers
         else:
             self.lstm = None
             self.att = None
+            self.channels = None
+            self.num_layers = None
 
         self.reset_parameters()
 
@@ -82,4 +86,7 @@ def forward(self, xs: List[Tensor]) -> Tensor:
             return (x * alpha.unsqueeze(-1)).sum(dim=1)
 
     def __repr__(self) -> str:
+        if self.mode == 'lstm':
+            return f'{self.__class__.__name__}({self.mode}, ' + \
+                   f'Channels: {self.channels}, Layers: {self.num_layers})'
         return f'{self.__class__.__name__}({self.mode})'


Code review comments:
diff_hunk: @@ -82,4 +86,7 @@ def forward(self, xs: List[Tensor]) -> Tensor:
             return (x * alpha.unsqueeze(-1)).sum(dim=1)
 
     def __repr__(self) -> str:
+        if self.mode == 'lstm':
+            return f'{self.__class__.__name__}({self.mode}, ' + \
+                   f'Channels: {self.channels}, Layers: {self.num_layers})'
path: torch_geometric/nn/models/jumping_knowledge.py
text: @user1:
```suggestion
            return (f'{self.__class__.__name__}({self.mode}, '
                    f'channels={self.channels}, layers={self.num_layers})')
```

diff_hunk: @@ -29,7 +29,8 @@ def test_jumping_knowledge():
         assert torch.allclose(jit(xs), out)
 
     model = JumpingKnowledge('lstm', channels, num_layers)
-    assert model.__repr__() == 'JumpingKnowledge(lstm)'
+    assert model.__repr__() == 'JumpingKnowledge(lstm, Channels: ' +\
+                               f'{channels}, Layers: {num_layers})'
path: test/nn/models/test_jumping_knowledge.py
text: @user1:
```suggestion
    assert model.__repr__() == (f'JumpingKnowledge(lstm, channels='
                                f'{channels}, layers={num_layers})')
```