code: def forward(
      self, fatoms, fbonds, agraph, bgraph, scope, tree_message
  ):
      fatoms = create_var(fatoms)
      fbonds = create_var(fbonds)
      agraph = create_var(agraph)
      bgraph = create_var(bgraph)
  
      binput = self.W_i(fbonds)
      graph_message = F.relu(binput)
  
      for i in range(self.depth - 1):
          message = torch.cat([tree_message, graph_message], dim=0)
          nei_message = index_select_ND(message, 0, bgraph)
          # assuming tree_message[0] == vec(0)
          nei_message = nei_message.sum(dim=1)
          nei_message = self.W_h(nei_message)
          graph_message = F.relu(binput + nei_message)
  
      message = torch.cat([tree_message, graph_message], dim=0)
      nei_message = index_select_ND(message, 0, agraph)
      nei_message = nei_message.sum(dim=1)
      ainput = torch.cat([fatoms, nei_message], dim=1)
      atom_hiddens = F.relu(self.W_o(ainput))
  
      mol_vecs = []
      for st, le in scope:
          mol_vec = atom_hiddens.narrow(0, st, le).sum(dim=0) / le
          mol_vecs.append(mol_vec)
  
      mol_vecs = torch.stack(mol_vecs, dim=0)
      return mol_vecs
  
file_path: DIG/dig/ggraph/method/JTVAE/fast_jtnn/jtmpn.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  import rdkit.Chem
  from .nnutils import create_var, index_select_ND
label_desc: Use image node_id 2 for calling the JTMPN obj's underlying member method code with example usage: obj.forward(fatoms, fbonds, agraph, bgraph, scope, tree_message) and returns: mol_vecs