code: def forward(self, cand_batch, tree_mess):
      fatoms, fbonds = [], []
      in_bonds, all_bonds = [], []
      mess_dict, all_mess = {}, [
          create_var(torch.zeros(self.hidden_size))
      ]  # Ensure index 0 is vec(0)
      total_atoms = 0
      scope = []
  
      for e, vec in tree_mess.items():
          mess_dict[e] = len(all_mess)
          all_mess.append(vec)
  
      for mol, all_nodes in cand_batch:
          n_atoms = mol.GetNumAtoms()
  
          for atom in mol.GetAtoms():
              fatoms.append(atom_features(atom))
              in_bonds.append([])
  
          for bond in mol.GetBonds():
              a1 = bond.GetBeginAtom()
              a2 = bond.GetEndAtom()
              x = a1.GetIdx() + total_atoms
              y = a2.GetIdx() + total_atoms
              # Here x_nid,y_nid could be 0
              x_nid, y_nid = a1.GetAtomMapNum(), a2.GetAtomMapNum()
              x_bid = all_nodes[x_nid - 1].idx if x_nid > 0 else -1
              y_bid = all_nodes[y_nid - 1].idx if y_nid > 0 else -1
  
              bfeature = bond_features(bond)
  
              # bond idx offseted by len(all_mess)
              b = len(all_mess) + len(all_bonds)
              all_bonds.append((x, y))
              fbonds.append(torch.cat([fatoms[x], bfeature], 0))
              in_bonds[y].append(b)
  
              b = len(all_mess) + len(all_bonds)
              all_bonds.append((y, x))
              fbonds.append(torch.cat([fatoms[y], bfeature], 0))
              in_bonds[x].append(b)
  
              if x_bid >= 0 and y_bid >= 0 and x_bid != y_bid:
                  if (x_bid, y_bid) in mess_dict:
                      mess_idx = mess_dict[(x_bid, y_bid)]
                      in_bonds[y].append(mess_idx)
                  if (y_bid, x_bid) in mess_dict:
                      mess_idx = mess_dict[(y_bid, x_bid)]
                      in_bonds[x].append(mess_idx)
  
          scope.append((total_atoms, n_atoms))
          total_atoms += n_atoms
  
      total_bonds = len(all_bonds)
      total_mess = len(all_mess)
      fatoms = torch.stack(fatoms, 0)
      fbonds = torch.stack(fbonds, 0)
      agraph = torch.zeros(total_atoms, MAX_NB).long()
      bgraph = torch.zeros(total_bonds, MAX_NB).long()
      tree_message = torch.stack(all_mess, dim=0)
  
      for a in range(total_atoms):
          for i, b in enumerate(in_bonds[a]):
              agraph[a, i] = b
  
      for b1 in range(total_bonds):
          x, y = all_bonds[b1]
          # b2 is offseted by len(all_mess)
          for i, b2 in enumerate(in_bonds[x]):
              if b2 < total_mess or all_bonds[b2 - total_mess][0] != y:
                  bgraph[b1, i] = b2
  
      fatoms = create_var(fatoms)
      fbonds = create_var(fbonds)
      agraph = create_var(agraph)
      bgraph = create_var(bgraph)
  
      binput = self.W_i(fbonds)
      graph_message = nn.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)
          nei_message = nei_message.sum(dim=1)
          nei_message = self.W_h(nei_message)
          graph_message = nn.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 = nn.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_bo.py
imports: import torch
  import torch.nn
  from .nnutils import create_var, index_select_ND
  import rdkit.Chem
label_desc: Use image node_id 2 for calling the JTMPNBO obj's underlying member method code with example usage: obj.forward(cand_batch, tree_mess) and returns: mol_vecs