class_docstr: Model for Graph Property Prediction.
  
  This model proceeds as follows:
  
  * Combine latest node representations and edge features in updating node representations,
      which involves multiple rounds of message passing
  * For each graph, compute its representation by combining the representations
      of all nodes in it, which involves a Set2Set layer.
  * Perform the final prediction using an MLP
  
  Examples
  --------
  
  >>> import deepchem as dc
  >>> import dgl
  >>> from deepchem.models.torch_models import MPNN
  >>> smiles = ["C1CCC1", "C1=CC=CN=C1"]
  >>> featurizer = dc.feat.MolGraphConvFeaturizer(use_edges=True)
  >>> graphs = featurizer.featurize(smiles)
  >>> print(type(graphs[0]))
  <class 'deepchem.feat.graph_data.GraphData'>
  >>> dgl_graphs = [graphs[i].to_dgl_graph(self_loop=True) for i in range(len(graphs))]
  >>> # Batch two graphs into a graph of two connected components
  >>> batch_dgl_graph = dgl.batch(dgl_graphs)
  >>> model = MPNN(n_tasks=1, mode='regression')
  >>> preds = model(batch_dgl_graph)
  >>> print(type(preds))
  <class 'torch.Tensor'>
  >>> preds.shape == (2, 1)
  True
  
  References
  ----------
  .. [1] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, George E. Dahl.
      "Neural Message Passing for Quantum Chemistry." ICML 2017.
  
  Notes
  -----
  This class requires DGL (https://github.com/dmlc/dgl) and DGL-LifeSci
  (https://github.com/awslabs/dgl-lifesci) to be installed.
docstr: Predict graph labels
  
  Parameters
  ----------
  g: DGLGraph
      A DGLGraph for a batch of graphs. It stores the node features in
      ``dgl_graph.ndata[self.nfeat_name]`` and edge features in
      ``dgl_graph.edata[self.efeat_name]``.
  
  Returns
  -------
  torch.Tensor
      The model output.
  
  * When self.mode = 'regression',
      its shape will be ``(dgl_graph.batch_size, self.n_tasks)``.
  * When self.mode = 'classification', the output consists of probabilities
      for classes. Its shape will be
      ``(dgl_graph.batch_size, self.n_tasks, self.n_classes)`` if self.n_tasks > 1;
      its shape will be ``(dgl_graph.batch_size, self.n_classes)`` if self.n_tasks is 1.
  torch.Tensor, optional
      This is only returned when self.mode = 'classification', the output consists of the
      logits for classes before softmax.
code: def forward(self, g):
      """Predict graph labels
  
      Parameters
      ----------
      g: DGLGraph
          A DGLGraph for a batch of graphs. It stores the node features in
          ``dgl_graph.ndata[self.nfeat_name]`` and edge features in
          ``dgl_graph.edata[self.efeat_name]``.
  
      Returns
      -------
      torch.Tensor
          The model output.
  
      * When self.mode = 'regression',
          its shape will be ``(dgl_graph.batch_size, self.n_tasks)``.
      * When self.mode = 'classification', the output consists of probabilities
          for classes. Its shape will be
          ``(dgl_graph.batch_size, self.n_tasks, self.n_classes)`` if self.n_tasks > 1;
          its shape will be ``(dgl_graph.batch_size, self.n_classes)`` if self.n_tasks is 1.
      torch.Tensor, optional
          This is only returned when self.mode = 'classification', the output consists of the
          logits for classes before softmax.
      """
      node_feats = g.ndata[self.nfeat_name]
      edge_feats = g.edata[self.efeat_name]
      out = self.model(g, node_feats, edge_feats)
  
      if self.mode == "classification":
          if self.n_tasks == 1:
              logits = out.view(-1, self.n_classes)
              softmax_dim = 1
          else:
              logits = out.view(-1, self.n_tasks, self.n_classes)
              softmax_dim = 2
          proba = F.softmax(logits, dim=softmax_dim)
          return proba, logits
      else:
          return out
  
file_path: deepchem/deepchem/models/torch_models/mpnn.py
imports: import torch.nn
  import torch.nn.functional
  from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
  from deepchem.models.torch_models.torch_model import TorchModel
label_desc: Use image node_id 2 for calling the MPNN obj's underlying member method code with example usage: obj.forward(g) and returns: proba,  logits, out