Neural Network Components
============================================

.. rubric:: Multi-Head Attention with Node Masking
.. automodule:: torch_molecule.nn.attention
   :members:
   :undoc-members:
   :show-inheritance:

.. rubric:: Embedding methods for different input features
.. automodule:: torch_molecule.nn.embedder
   :members:
   :undoc-members:
   :show-inheritance:

.. rubric:: Layers for Graph Neural Networks
.. automodule:: torch_molecule.nn.gnn
   :members:
   :undoc-members:
   :show-inheritance:

.. rubric:: Multi-Layer Perceptron
.. automodule:: torch_molecule.nn.mlp
   :members:
   :undoc-members:
   :show-inheritance: