code: def forward(self, mol_batch, x_tree_vecs):
      pred_hiddens, pred_contexts, pred_targets = [], [], []
      stop_hiddens, stop_contexts, stop_targets = [], [], []
      traces = []
      for mol_tree in mol_batch:
          s = []
          dfs(s, mol_tree.nodes[0], -1)
          traces.append(s)
          for node in mol_tree.nodes:
              node.neighbors = []
  
      # Predict Root
      batch_size = len(mol_batch)
      pred_hiddens.append(
          create_var(torch.zeros(len(mol_batch), self.hidden_size))
      )
      pred_targets.extend(
          [mol_tree.nodes[0].wid for mol_tree in mol_batch]
      )
      pred_contexts.append(
          create_var(torch.LongTensor(range(batch_size)))
      )
  
      max_iter = max([len(tr) for tr in traces])
      padding = create_var(torch.zeros(self.hidden_size), False)
      h = {}
  
      for t in range(max_iter):
          prop_list = []
          batch_list = []
          for i, plist in enumerate(traces):
              if t < len(plist):
                  prop_list.append(plist[t])
                  batch_list.append(i)
  
          cur_x = []
          cur_h_nei, cur_o_nei = [], []
  
          for node_x, real_y, _ in prop_list:
              # Neighbors for message passing (target not included)
              cur_nei = [
                  h[(node_y.idx, node_x.idx)]
                  for node_y in node_x.neighbors
                  if node_y.idx != real_y.idx
              ]
              pad_len = MAX_NB - len(cur_nei)
              cur_h_nei.extend(cur_nei)
              cur_h_nei.extend([padding] * pad_len)
  
              # Neighbors for stop prediction (all neighbors)
              cur_nei = [
                  h[(node_y.idx, node_x.idx)]
                  for node_y in node_x.neighbors
              ]
              pad_len = MAX_NB - len(cur_nei)
              cur_o_nei.extend(cur_nei)
              cur_o_nei.extend([padding] * pad_len)
  
              # Current clique embedding
              cur_x.append(node_x.wid)
  
          # Clique embedding
          cur_x = create_var(torch.LongTensor(cur_x))
          cur_x = self.embedding(cur_x)
  
          # Message passing
          cur_h_nei = torch.stack(cur_h_nei, dim=0).view(
              -1, MAX_NB, self.hidden_size
          )
          new_h = GRU(
              cur_x, cur_h_nei, self.W_z, self.W_r, self.U_r, self.W_h
          )
  
          # Node Aggregate
          cur_o_nei = torch.stack(cur_o_nei, dim=0).view(
              -1, MAX_NB, self.hidden_size
          )
          cur_o = cur_o_nei.sum(dim=1)
  
          # Gather targets
          pred_target, pred_list = [], []
          stop_target = []
          for i, m in enumerate(prop_list):
              node_x, node_y, direction = m
              x, y = node_x.idx, node_y.idx
              h[(x, y)] = new_h[i]
              node_y.neighbors.append(node_x)
              if direction == 1:
                  pred_target.append(node_y.wid)
                  pred_list.append(i)
              stop_target.append(direction)
  
          # Hidden states for stop prediction
          cur_batch = create_var(torch.LongTensor(batch_list))
          stop_hidden = torch.cat([cur_x, cur_o], dim=1)
          stop_hiddens.append(stop_hidden)
          stop_contexts.append(cur_batch)
          stop_targets.extend(stop_target)
  
          # Hidden states for clique prediction
          if len(pred_list) > 0:
              batch_list = [batch_list[i] for i in pred_list]
              cur_batch = create_var(torch.LongTensor(batch_list))
              pred_contexts.append(cur_batch)
  
              cur_pred = create_var(torch.LongTensor(pred_list))
              pred_hiddens.append(new_h.index_select(0, cur_pred))
              pred_targets.extend(pred_target)
  
      # Last stop at root
      cur_x, cur_o_nei = [], []
      for mol_tree in mol_batch:
          node_x = mol_tree.nodes[0]
          cur_x.append(node_x.wid)
          cur_nei = [
              h[(node_y.idx, node_x.idx)] for node_y in node_x.neighbors
          ]
          pad_len = MAX_NB - len(cur_nei)
          cur_o_nei.extend(cur_nei)
          cur_o_nei.extend([padding] * pad_len)
  
      cur_x = create_var(torch.LongTensor(cur_x))
      cur_x = self.embedding(cur_x)
      cur_o_nei = torch.stack(cur_o_nei, dim=0).view(
          -1, MAX_NB, self.hidden_size
      )
      cur_o = cur_o_nei.sum(dim=1)
  
      stop_hidden = torch.cat([cur_x, cur_o], dim=1)
      stop_hiddens.append(stop_hidden)
      stop_contexts.append(
          create_var(torch.LongTensor(range(batch_size)))
      )
      stop_targets.extend([0] * len(mol_batch))
  
      # Predict next clique
      pred_contexts = torch.cat(pred_contexts, dim=0)
      pred_hiddens = torch.cat(pred_hiddens, dim=0)
      pred_scores = self.aggregate(
          pred_hiddens, pred_contexts, x_tree_vecs, "word"
      )
      pred_targets = create_var(torch.LongTensor(pred_targets))
  
      pred_loss = self.pred_loss(pred_scores, pred_targets) / len(
          mol_batch
      )
      _, preds = torch.max(pred_scores, dim=1)
      pred_acc = torch.eq(preds, pred_targets).float()
      pred_acc = torch.sum(pred_acc) / pred_targets.nelement()
  
      # Predict stop
      stop_contexts = torch.cat(stop_contexts, dim=0)
      stop_hiddens = torch.cat(stop_hiddens, dim=0)
      stop_hiddens = F.relu(self.U_i(stop_hiddens))
      stop_scores = self.aggregate(
          stop_hiddens, stop_contexts, x_tree_vecs, "stop"
      )
      stop_scores = stop_scores.squeeze(-1)
      stop_targets = create_var(torch.Tensor(stop_targets))
  
      stop_loss = self.stop_loss(stop_scores, stop_targets) / len(
          mol_batch
      )
      stops = torch.ge(stop_scores, 0).float()
      stop_acc = torch.eq(stops, stop_targets).float()
      stop_acc = torch.sum(stop_acc) / stop_targets.nelement()
  
      return pred_loss, stop_loss, pred_acc.item(), stop_acc.item()
  
file_path: DIG/dig/ggraph/method/JTVAE/fast_jtnn/jtnn_dec.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .mol_tree import MolTree, MolTreeNode
  from .chemutils import enum_assemble, set_atommap
  from .nnutils import create_var, GRU
label_desc: Use image node_id 3 for calling the JTNNDecoder obj's underlying member method code with example usage: obj.forward(mol_batch, x_tree_vecs) and returns: pred_loss,  stop_loss,  pred_acc,  stop_acc