chemprop.nn.message_passing.proto
=================================

.. py:module:: chemprop.nn.message_passing.proto


Classes
-------

.. autoapisummary::

   chemprop.nn.message_passing.proto.MessagePassing
   chemprop.nn.message_passing.proto.MABMessagePassing


Module Contents
---------------

.. py:class:: MessagePassing(*args, **kwargs)

   Bases: :py:obj:`torch.nn.Module`, :py:obj:`chemprop.nn.hparams.HasHParams`


   A :class:`MessagePassing` module encodes a batch of molecular graphs
   using message passing to learn vertex-level hidden representations.


   .. py:attribute:: output_dim
      :type:  int


   .. py:method:: forward(bmg, V_d = None)
      :abstractmethod:


      Encode a batch of molecular graphs.

      :param bmg: the batch of :class:`~chemprop.featurizers.molgraph.MolGraph`\s to encode
      :type bmg: BatchMolGraph
      :param V_d: an optional tensor of shape `V x d_vd` containing additional descriptors for each vertex
                  in the batch. These will be concatenated to the learned vertex descriptors and
                  transformed before the readout phase.
      :type V_d: Tensor | None, default=None

      :returns: a tensor of shape `V x d_h` or `V x (d_h + d_vd)` containing the hidden representation
                of each vertex in the batch of graphs. The feature dimension depends on whether
                additional vertex descriptors were provided
      :rtype: Tensor



.. py:class:: MABMessagePassing(*args, **kwargs)

   Bases: :py:obj:`torch.nn.Module`, :py:obj:`chemprop.nn.hparams.HasHParams`


   A :class:`MABMessagePassing` module encodes a batch of molecular graphs
   using message passing to learn both vertex-level and edge-level hidden representations.


   .. py:attribute:: output_dims
      :type:  tuple[int | None, int | None]


   .. py:method:: forward(bmg, V_d = None, E_d = None)
      :abstractmethod:


      Encode a batch of molecular graphs.

      :param bmg: the batch of :class:`~chemprop.featurizers.molgraph.MolGraph`\s to encode
      :type bmg: BatchMolGraph
      :param V_d: an optional tensor of shape `V x d_vd` containing additional descriptors for each vertex
                  in the batch. These will be concatenated to the learned vertex descriptors and
                  transformed before the readout phase.
      :type V_d: Tensor | None, default=None
      :param E_d: an optional tensor of shape `E x d_ed` containing additional descriptors for each
                  directed edge in the batch. These will be concatenated to the learned edge descriptors
                  and transformed before the readout phase. NOTE: There are two directed edges per graph
                  connection. If the extra descriptors are for the connections, each row should be
                  repeated twice in the tensor, once for each direction, potentially using
                  ``E_d = np.repeat(E_d, repeats=2, axis=0)``.
      :type E_d: Tensor | None, default=None

      :returns: Two tensors of shape `V x d_h` or `V x (d_h + d_vd)` and `E x dh` or `E x (dh + d_ed)`
                containing the hidden representation of each vertex and edge in the batch of graphs.
                The feature dimension depends on whether additional atom/bond descriptors were provided.
                If either the vertex or edge hidden representations are not needed, computing the
                corresponding tensor can be suppresed by setting either return_vertex_embeddings or
                return_edge_embeddings to `False` when initializing the module.
      :rtype: tuple[Tensor | None, Tensor | None]