Repository: pyg-team/pytorch_geometric

Issue:
Improving documentation for Set2Set layer
### 📚 Describe the documentation issue

I am new to `pytorch_geometric` ecosystem and I was exploring it. At the first glance to the `Set2Set` layer in the [docs](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.glob.Set2Set), it is not clear what the inputs `x` and `batch` are to the forward pass.

If I am not wrong, `x` represents the node features of the graph and `batch` represents a mapping between the node features to their graph identifiers.


### Suggest a potential alternative/fix

I was wondering whether it will be good to include it to the docs or maybe also add typing.
Potential fix in `nn.glob.set2set.py`:
```
def forward(self, x: torch.Tensor, batch: torch.Tensor):
        r"""
        Args:
             x: The input node features.
             batch: A one dimension tensor representing a mapping between nodes and their graphs
        """
```


Discussion:
@arunppsg I think its a good idea to add type hints. Please feel free to contribute this.
There is also some recent effort to bring shape information to all GNN operators, see [here](https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/gin_conv.py#L41-L48). We can likely also do this for any other PyG operator.
Putting up shape information will be super useful. I will put up a small PR regarding the suggestions later this week. Thanks!

Resolving patch:
diff --git a/torch_geometric/nn/glob/set2set.py b/torch_geometric/nn/glob/set2set.py
--- a/torch_geometric/nn/glob/set2set.py
+++ b/torch_geometric/nn/glob/set2set.py
@@ -1,5 +1,9 @@
+from typing import Optional
+
 import torch
+from torch import Tensor
 from torch_scatter import scatter_add
+
 from torch_geometric.utils import softmax
 
 
@@ -27,8 +31,17 @@ class Set2Set(torch.nn.Module):
             :obj:`num_layers=2` would mean stacking two LSTMs together to form
             a stacked LSTM, with the second LSTM taking in outputs of the first
             LSTM and computing the final results. (default: :obj:`1`)
+
+    Shapes:
+        - **input:**
+          node features :math:`(|\mathcal{V}|, F)`,
+          batch vector :math:`(|\mathcal{V}|)` *(optional)*
+        - **output:**
+          set features :math:`(|\mathcal{G}|, 2 * F)` where
+          :math:`|\mathcal{G}|` denotes the number of graphs in the batch
     """
-    def __init__(self, in_channels, processing_steps, num_layers=1):
+    def __init__(self, in_channels: int, processing_steps: int,
+                 num_layers: int = 1):
         super().__init__()
 
         self.in_channels = in_channels
@@ -44,8 +57,16 @@ def __init__(self, in_channels, processing_steps, num_layers=1):
     def reset_parameters(self):
         self.lstm.reset_parameters()
 
-    def forward(self, x, batch):
-        """"""
+    def forward(self, x: Tensor, batch: Optional[Tensor] = None) -> Tensor:
+        r"""
+        Args:
+            x (Tensor): The input node features.
+            batch (LongTensor, optional): A vector that maps each node to its
+                respective graph identifier. (default: :obj:`None`)
+        """
+        if batch is None:
+            batch = x.new_zeros(x.size(0), dtype=torch.int64)
+
         batch_size = batch.max().item() + 1
 
         h = (x.new_zeros((self.num_layers, batch_size, self.in_channels)),


Test patch: