class_docstr: Directed Message Passing Neural Network
  
  In this class, we define the various encoder layers and establish a sequential model for the Directed Message Passing Neural Network (D-MPNN) [1]_.
  We also define the forward call of this model in the forward function.
  
  Example
  -------
  >>> import deepchem as dc
  >>> from torch_geometric.data import Data, Batch
  >>> # Get data
  >>> input_smile = "CC"
  >>> feat = dc.feat.DMPNNFeaturizer(features_generators=['morgan'])
  >>> graph = feat.featurize(input_smile)
  >>> mapper = _MapperDMPNN(graph[0])
  >>> atom_features, f_ini_atoms_bonds, atom_to_incoming_bonds, mapping, global_features = mapper.values
  >>> atom_features = torch.from_numpy(atom_features).float()
  >>> f_ini_atoms_bonds = torch.from_numpy(f_ini_atoms_bonds).float()
  >>> atom_to_incoming_bonds = torch.from_numpy(atom_to_incoming_bonds)
  >>> mapping = torch.from_numpy(mapping)
  >>> global_features = torch.from_numpy(global_features).float()
  >>> data = [Data(atom_features=atom_features,    f_ini_atoms_bonds=f_ini_atoms_bonds,    atom_to_incoming_bonds=atom_to_incoming_bonds,    mapping=mapping, global_features=global_features)]
  >>> # Prepare batch (size 1)
  >>> pyg_batch = Batch()
  >>> pyg_batch = pyg_batch.from_data_list(data)
  >>> # Initialize the model
  >>> model = DMPNN(mode='regression', global_features_size=2048, n_tasks=2)
  >>> # Get the forward call of the model for this batch.
  >>> output = model(pyg_batch)
  
  References
  ----------
  .. [1] Analyzing Learned Molecular Representations for Property Prediction https://arxiv.org/pdf/1904.01561.pdf
class_docstr_tok: ["Directed","Message","Passing","Neural","Network","In","this","class",",","we","define","the","various","encoder","layers","and","establish","a","sequential","model","for","the","Directed","Message","Passing","Neural","Network","(","D-MPNN",")","[","1","]","_",".","We","also","define","the","forward","call","of","this","model","in","the","forward","function",".","Example","--","--","--","-",">",">",">","import","deepchem","as","dc",">",">",">","from","torch_geometric.data","import","Data",",","Batch",">",">",">","#","Get","data",">",">",">","input_smile","=","``","CC","''",">",">",">","feat","=","dc.feat.DMPNNFeaturizer","(","features_generators=","[","'morgan","'","]",")",">",">",">","graph","=","feat.featurize","(","input_smile",")",">",">",">","mapper","=","_MapperDMPNN","(","graph","[","0","]",")",">",">",">","atom_features",",","f_ini_atoms_bonds",",","atom_to_incoming_bonds",",","mapping",",","global_features","=","mapper.values",">",">",">","atom_features","=","torch.from_numpy","(","atom_features",")",".float","(",")",">",">",">","f_ini_atoms_bonds","=","torch.from_numpy","(","f_ini_atoms_bonds",")",".float","(",")",">",">",">","atom_to_incoming_bonds","=","torch.from_numpy","(","atom_to_incoming_bonds",")",">",">",">","mapping","=","torch.from_numpy","(","mapping",")",">",">",">","global_features","=","torch.from_numpy","(","global_features",")",".float","(",")",">",">",">","data","=","[","Data","(","atom_features=atom_features",",","f_ini_atoms_bonds=f_ini_atoms_bonds",",","atom_to_incoming_bonds=atom_to_incoming_bonds",",","mapping=mapping",",","global_features=global_features",")","]",">",">",">","#","Prepare","batch","(","size","1",")",">",">",">","pyg_batch","=","Batch","(",")",">",">",">","pyg_batch","=","pyg_batch.from_data_list","(","data",")",">",">",">","#","Initialize","the","model",">",">",">","model","=","DMPNN","(","mode='regression","'",",","global_features_size=2048",",","n_tasks=2",")",">",">",">","#","Get","the","forward","call","of","the","model","for","this","batch",".",">",">",">","output","=","model","(","pyg_batch",")","References","--","--","--","--","--","..","[","1","]","Analyzing","Learned","Molecular","Representations","for","Property","Prediction","https",":","\/\/arxiv.org\/pdf\/1904.01561.pdf"]
docstr: Parameters
  ----------
  data: Batch
      A pytorch-geometric batch containing tensors for:
  
      - atom_features
      - f_ini_atoms_bonds
      - atom_to_incoming_bonds
      - mapping
      - global_features
  
  The `molecules_unbatch_key` is also derived from the batch.
  (List containing number of atoms in various molecules of the batch)
  
  Returns
  -------
  output: Union[torch.Tensor, Sequence[torch.Tensor]]
      Predictions for the graphs
docstr_tok: ["Parameters","--","--","--","--","--","data",":","Batch","A","pytorch-geometric","batch","containing","tensors","for",":","-","atom_features","-","f_ini_atoms_bonds","-","atom_to_incoming_bonds","-","mapping","-","global_features","The","`","molecules_unbatch_key","`","is","also","derived","from","the","batch",".","(","List","containing","number","of","atoms","in","various","molecules","of","the","batch",")","Returns","--","--","--","-","output",":","Union","[","torch.Tensor",",","Sequence","[","torch.Tensor","]","]","Predictions","for","the","graphs"]
code: def forward(
      self, pyg_batch: Batch
  ) -> Union[torch.Tensor, Sequence[torch.Tensor]]:
      """
      Parameters
      ----------
      data: Batch
          A pytorch-geometric batch containing tensors for:
  
          - atom_features
          - f_ini_atoms_bonds
          - atom_to_incoming_bonds
          - mapping
          - global_features
  
      The `molecules_unbatch_key` is also derived from the batch.
      (List containing number of atoms in various molecules of the batch)
  
      Returns
      -------
      output: Union[torch.Tensor, Sequence[torch.Tensor]]
          Predictions for the graphs
      """
      atom_features: torch.Tensor = pyg_batch["atom_features"]
      f_ini_atoms_bonds: torch.Tensor = pyg_batch["f_ini_atoms_bonds"]
      atom_to_incoming_bonds: torch.Tensor = pyg_batch[
          "atom_to_incoming_bonds"
      ]
      mapping: torch.Tensor = pyg_batch["mapping"]
      global_features: torch.Tensor = pyg_batch["global_features"]
  
      # Steps to get `molecules_unbatch_key`:
      # 1. Get the tensor containing the indices of first atoms of each molecule
      # 2. Get the tensor containing number of atoms of each molecule
      #     by taking the difference between consecutive indices.
      # 3. Convert the tensor to a list.
      molecules_unbatch_key: List = torch.diff(
          pyg_batch._slice_dict["atom_features"]
      ).tolist()
  
      # num_molecules x (enc_hidden + global_features_size)
      encodings: torch.Tensor = self.encoder(
          atom_features,
          f_ini_atoms_bonds,
          atom_to_incoming_bonds,
          mapping,
          global_features,
          molecules_unbatch_key,
      )
  
      # ffn_output (`self.n_tasks` or `self.n_tasks * self.n_classes`)
      output: torch.Tensor = self.ffn(encodings)
  
      final_output: Union[torch.Tensor, Sequence[torch.Tensor]]
  
      if self.mode == "regression":
          final_output = output
      elif self.mode == "classification":
          if self.n_tasks == 1:
              output = output.view(-1, self.n_classes)
              final_output = (
                  nn.functional.softmax(output, dim=1),
                  output,
              )
          else:
              output = output.view(-1, self.n_tasks, self.n_classes)
              final_output = (
                  nn.functional.softmax(output, dim=2),
                  output,
              )
  
      return final_output
  
imports: import numpy
  import torch
  import torch.nn
  from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
  from deepchem.models.torch_models import layers
  from deepchem.models.torch_models import TorchModel
  from torch_geometric.data import Data, Batch
  from deepchem.feat import GraphData
  from deepchem.data import Dataset
  from typing import Union, List, Sequence, Optional, Iterable, Tuple
label_desc: Use image node_id 2 for calling the DMPNN obj's underlying member method code with example usage: obj.forward(pyg_batch) and returns: final_output