Add torch_geometric.nn.attention to docs

docs/source/modules/nn.rst

@@ -157,6 +157,20 @@ Finally, we added full support for customization of aggregations into the :class
      {{ name }}
    {% endfor %}
 
+Attention
+---------
+
+.. currentmodule:: torch_geometric.nn.attention
+
+.. autosummary::
+   :nosignatures:
+   :toctree: ../generated
+   :template: autosummary/nn.rst
+
+   {% for name in torch_geometric.nn.attention.classes %}
+     {{ name }}
+   {% endfor %}
+
 Normalization Layers
 --------------------
 
@@ -226,17 +240,30 @@ KGE Models
 Encodings
 ---------
 
-.. automodule:: torch_geometric.nn.encoding
-   :members:
-   :undoc-members:
-   :exclude-members: training
+.. currentmodule:: torch_geometric.nn.encoding
+
+.. autosummary::
+   :nosignatures:
+   :toctree: ../generated
+   :template: autosummary/nn.rst
+
+   {% for name in torch_geometric.nn.encoding.classes %}
+     {{ name }}
+   {% endfor %}
 
 NLP
 ---------
 
-.. automodule:: torch_geometric.nn.nlp
-   :members:
-   :undoc-members:
+.. currentmodule:: torch_geometric.nn.nlp
+
+.. autosummary::
+   :nosignatures:
+   :toctree: ../generated
+   :template: autosummary/nn.rst
+
+   {% for name in torch_geometric.nn.nlp.classes %}
+     {{ name }}
+   {% endfor %}
 
 Functional
 ----------

torch_geometric/nn/__init__.py

@@ -18,6 +18,7 @@
 from .models import *  # noqa
 from .functional import *  # noqa
 from .nlp import *  # noqa
+from .attention import *  # noqa
 
 __all__ = [
     'Reshape',

torch_geometric/nn/attention/__init__.py

@@ -2,7 +2,7 @@
 from .qformer import QFormer
 from .sgformer import SGFormerAttention
 
-__all__ = [
+__all__ = classes = [
     'PerformerAttention',
     'QFormer',
     'SGFormerAttention',

torch_geometric/nn/encoding.py

@@ -3,6 +3,11 @@
 import torch
 from torch import Tensor
 
+__all__ = classes = [
+    'PositionalEncoding',
+    'TemporalEncoding',
+]
+
 
 class PositionalEncoding(torch.nn.Module):
     r"""The positional encoding scheme from the `"Attention Is All You Need"

torch_geometric/nn/nlp/sentence_transformer.py

@@ -14,6 +14,13 @@ class PoolingStrategy(Enum):
 
 
 class SentenceTransformer(torch.nn.Module):
+    r"""A wrapper around a Sentence-Transformer from HuggingFace.
+
+    Args:
+        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"BERT"`.
+        pooling_strategy (str, optional): The pooling strategy to use
+            for generating node embeddings. (default: :obj:`"mean"`)
+    """
     def __init__(
         self,
         model_name: str,

torch_geometric/nn/nlp/vision_transformer.py

@@ -5,6 +5,11 @@
 
 
 class VisionTransformer(torch.nn.Module):
+    r"""A wrapper around a Vision-Transformer from HuggingFace.
+
+    Args:
+        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"ViT"`.
+    """
     def __init__(
         self,
         model_name: str,