prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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
          """
  ```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `torch_geometric/nn/glob/set2set.py`
  Content:
  ```
  1 import torch
  2 from torch_scatter import scatter_add
  3 from torch_geometric.utils import softmax
  4 
  5 
  6 class Set2Set(torch.nn.Module):
  7     r"""The global pooling operator based on iterative content-based attention
  8     from the `"Order Matters: Sequence to sequence for sets"
  9     <https://arxiv.org/abs/1511.06391>`_ paper
  10 
  11     .. math::
  12         \mathbf{q}_t &= \mathrm{LSTM}(\mathbf{q}^{*}_{t-1})
  13 
  14         \alpha_{i,t} &= \mathrm{softmax}(\mathbf{x}_i \cdot \mathbf{q}_t)
  15 
  16         \mathbf{r}_t &= \sum_{i=1}^N \alpha_{i,t} \mathbf{x}_i
  17 
  18         \mathbf{q}^{*}_t &= \mathbf{q}_t \, \Vert \, \mathbf{r}_t,
  19 
  20     where :math:`\mathbf{q}^{*}_T` defines the output of the layer with twice
  21     the dimensionality as the input.
  22 
  23     Args:
  24         in_channels (int): Size of each input sample.
  25         processing_steps (int): Number of iterations :math:`T`.
  26         num_layers (int, optional): Number of recurrent layers, *.e.g*, setting
  27             :obj:`num_layers=2` would mean stacking two LSTMs together to form
  28             a stacked LSTM, with the second LSTM taking in outputs of the first
  29             LSTM and computing the final results. (default: :obj:`1`)
  30     """
  31     def __init__(self, in_channels, processing_steps, num_layers=1):
  32         super().__init__()
  33 
  34         self.in_channels = in_channels
  35         self.out_channels = 2 * in_channels
  36         self.processing_steps = processing_steps
  37         self.num_layers = num_layers
  38 
  39         self.lstm = torch.nn.LSTM(self.out_channels, self.in_channels,
  40                                   num_layers)
  41 
  42         self.reset_parameters()
  43 
  44     def reset_parameters(self):
  45         self.lstm.reset_parameters()
  46 
  47     def forward(self, x, batch):
  48         """"""
  49         batch_size = batch.max().item() + 1
  50 
  51         h = (x.new_zeros((self.num_layers, batch_size, self.in_channels)),
  52              x.new_zeros((self.num_layers, batch_size, self.in_channels)))
  53         q_star = x.new_zeros(batch_size, self.out_channels)
  54 
  55         for _ in range(self.processing_steps):
  56             q, h = self.lstm(q_star.unsqueeze(0), h)
  57             q = q.view(batch_size, self.in_channels)
  58             e = (x * q.index_select(0, batch)).sum(dim=-1, keepdim=True)
  59             a = softmax(e, batch, num_nodes=batch_size)
  60             r = scatter_add(a * x, batch, dim=0, dim_size=batch_size)
  61             q_star = torch.cat([q, r], dim=-1)
  62 
  63         return q_star
  64 
  65     def __repr__(self) -> str:
  66         return (f'{self.__class__.__name__}({self.in_channels}, '
  67                 f'{self.out_channels})')
  68
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
               :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 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)),
  
verification_info: {"golden_diff": "diff --git a/torch_geometric/nn/glob/set2set.py b/torch_geometric/nn/glob/set2set.py\n--- a/torch_geometric/nn/glob/set2set.py\n+++ b/torch_geometric/nn/glob/set2set.py\n@@ -1,5 +1,9 @@\n+from typing import Optional\n+\n import torch\n+from torch import Tensor\n from torch_scatter import scatter_add\n+\n from torch_geometric.utils import softmax\n \n \n@@ -27,8 +31,17 @@\n             :obj:`num_layers=2` would mean stacking two LSTMs together to form\n             a stacked LSTM, with the second LSTM taking in outputs of the first\n             LSTM and computing the final results. (default: :obj:`1`)\n+\n+    Shapes:\n+        - **input:**\n+          node features :math:`(|\\mathcal{V}|, F)`,\n+          batch vector :math:`(|\\mathcal{V}|)` *(optional)*\n+        - **output:**\n+          set features :math:`(|\\mathcal{G}|, 2 * F)` where\n+          :math:`|\\mathcal{G}|` denotes the number of graphs in the batch\n     \"\"\"\n-    def __init__(self, in_channels, processing_steps, num_layers=1):\n+    def __init__(self, in_channels: int, processing_steps: int,\n+                 num_layers: int = 1):\n         super().__init__()\n \n         self.in_channels = in_channels\n@@ -44,8 +57,16 @@\n     def reset_parameters(self):\n         self.lstm.reset_parameters()\n \n-    def forward(self, x, batch):\n-        \"\"\"\"\"\"\n+    def forward(self, x: Tensor, batch: Optional[Tensor] = None) -> Tensor:\n+        r\"\"\"\n+        Args:\n+            x (Tensor): The input node features.\n+            batch (LongTensor, optional): A vector that maps each node to its\n+                respective graph identifier. (default: :obj:`None`)\n+        \"\"\"\n+        if batch is None:\n+            batch = x.new_zeros(x.size(0), dtype=torch.int64)\n+\n         batch_size = batch.max().item() + 1\n \n         h = (x.new_zeros((self.num_layers, batch_size, self.in_channels)),\n", "issue": "Improving documentation for Set2Set layer\n### \ud83d\udcda Describe the documentation issue\n\nI 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.\r\n\r\nIf 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.\r\n\n\n### Suggest a potential alternative/fix\n\nI was wondering whether it will be good to include it to the docs or maybe also add typing.\r\nPotential fix in `nn.glob.set2set.py`:\r\n```\r\ndef forward(self, x: torch.Tensor, batch: torch.Tensor):\r\n        r\"\"\"\r\n        Args:\r\n             x: The input node features.\r\n             batch: A one dimension tensor representing a mapping between nodes and their graphs\r\n        \"\"\"\r\n```\n", "before_files": [{"content": "import torch\nfrom torch_scatter import scatter_add\nfrom torch_geometric.utils import softmax\n\n\nclass Set2Set(torch.nn.Module):\n    r\"\"\"The global pooling operator based on iterative content-based attention\n    from the `\"Order Matters: Sequence to sequence for sets\"\n    <https://arxiv.org/abs/1511.06391>`_ paper\n\n    .. math::\n        \\mathbf{q}_t &= \\mathrm{LSTM}(\\mathbf{q}^{*}_{t-1})\n\n        \\alpha_{i,t} &= \\mathrm{softmax}(\\mathbf{x}_i \\cdot \\mathbf{q}_t)\n\n        \\mathbf{r}_t &= \\sum_{i=1}^N \\alpha_{i,t} \\mathbf{x}_i\n\n        \\mathbf{q}^{*}_t &= \\mathbf{q}_t \\, \\Vert \\, \\mathbf{r}_t,\n\n    where :math:`\\mathbf{q}^{*}_T` defines the output of the layer with twice\n    the dimensionality as the input.\n\n    Args:\n        in_channels (int): Size of each input sample.\n        processing_steps (int): Number of iterations :math:`T`.\n        num_layers (int, optional): Number of recurrent layers, *.e.g*, setting\n            :obj:`num_layers=2` would mean stacking two LSTMs together to form\n            a stacked LSTM, with the second LSTM taking in outputs of the first\n            LSTM and computing the final results. (default: :obj:`1`)\n    \"\"\"\n    def __init__(self, in_channels, processing_steps, num_layers=1):\n        super().__init__()\n\n        self.in_channels = in_channels\n        self.out_channels = 2 * in_channels\n        self.processing_steps = processing_steps\n        self.num_layers = num_layers\n\n        self.lstm = torch.nn.LSTM(self.out_channels, self.in_channels,\n                                  num_layers)\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        self.lstm.reset_parameters()\n\n    def forward(self, x, batch):\n        \"\"\"\"\"\"\n        batch_size = batch.max().item() + 1\n\n        h = (x.new_zeros((self.num_layers, batch_size, self.in_channels)),\n             x.new_zeros((self.num_layers, batch_size, self.in_channels)))\n        q_star = x.new_zeros(batch_size, self.out_channels)\n\n        for _ in range(self.processing_steps):\n            q, h = self.lstm(q_star.unsqueeze(0), h)\n            q = q.view(batch_size, self.in_channels)\n            e = (x * q.index_select(0, batch)).sum(dim=-1, keepdim=True)\n            a = softmax(e, batch, num_nodes=batch_size)\n            r = scatter_add(a * x, batch, dim=0, dim_size=batch_size)\n            q_star = torch.cat([q, r], dim=-1)\n\n        return q_star\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels})')\n", "path": "torch_geometric/nn/glob/set2set.py"}], "after_files": [{"content": "from typing import Optional\n\nimport torch\nfrom torch import Tensor\nfrom torch_scatter import scatter_add\n\nfrom torch_geometric.utils import softmax\n\n\nclass Set2Set(torch.nn.Module):\n    r\"\"\"The global pooling operator based on iterative content-based attention\n    from the `\"Order Matters: Sequence to sequence for sets\"\n    <https://arxiv.org/abs/1511.06391>`_ paper\n\n    .. math::\n        \\mathbf{q}_t &= \\mathrm{LSTM}(\\mathbf{q}^{*}_{t-1})\n\n        \\alpha_{i,t} &= \\mathrm{softmax}(\\mathbf{x}_i \\cdot \\mathbf{q}_t)\n\n        \\mathbf{r}_t &= \\sum_{i=1}^N \\alpha_{i,t} \\mathbf{x}_i\n\n        \\mathbf{q}^{*}_t &= \\mathbf{q}_t \\, \\Vert \\, \\mathbf{r}_t,\n\n    where :math:`\\mathbf{q}^{*}_T` defines the output of the layer with twice\n    the dimensionality as the input.\n\n    Args:\n        in_channels (int): Size of each input sample.\n        processing_steps (int): Number of iterations :math:`T`.\n        num_layers (int, optional): Number of recurrent layers, *.e.g*, setting\n            :obj:`num_layers=2` would mean stacking two LSTMs together to form\n            a stacked LSTM, with the second LSTM taking in outputs of the first\n            LSTM and computing the final results. (default: :obj:`1`)\n\n    Shapes:\n        - **input:**\n          node features :math:`(|\\mathcal{V}|, F)`,\n          batch vector :math:`(|\\mathcal{V}|)` *(optional)*\n        - **output:**\n          set features :math:`(|\\mathcal{G}|, 2 * F)` where\n          :math:`|\\mathcal{G}|` denotes the number of graphs in the batch\n    \"\"\"\n    def __init__(self, in_channels: int, processing_steps: int,\n                 num_layers: int = 1):\n        super().__init__()\n\n        self.in_channels = in_channels\n        self.out_channels = 2 * in_channels\n        self.processing_steps = processing_steps\n        self.num_layers = num_layers\n\n        self.lstm = torch.nn.LSTM(self.out_channels, self.in_channels,\n                                  num_layers)\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        self.lstm.reset_parameters()\n\n    def forward(self, x: Tensor, batch: Optional[Tensor] = None) -> Tensor:\n        r\"\"\"\n        Args:\n            x (Tensor): The input node features.\n            batch (LongTensor, optional): A vector that maps each node to its\n                respective graph identifier. (default: :obj:`None`)\n        \"\"\"\n        if batch is None:\n            batch = x.new_zeros(x.size(0), dtype=torch.int64)\n\n        batch_size = batch.max().item() + 1\n\n        h = (x.new_zeros((self.num_layers, batch_size, self.in_channels)),\n             x.new_zeros((self.num_layers, batch_size, self.in_channels)))\n        q_star = x.new_zeros(batch_size, self.out_channels)\n\n        for _ in range(self.processing_steps):\n            q, h = self.lstm(q_star.unsqueeze(0), h)\n            q = q.view(batch_size, self.in_channels)\n            e = (x * q.index_select(0, batch)).sum(dim=-1, keepdim=True)\n            a = softmax(e, batch, num_nodes=batch_size)\n            r = scatter_add(a * x, batch, dim=0, dim_size=batch_size)\n            q_star = torch.cat([q, r], dim=-1)\n\n        return q_star\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}({self.in_channels}, '\n                f'{self.out_channels})')\n", "path": "torch_geometric/nn/glob/set2set.py"}]}