code: def decode(self, mol_vec):
      assert mol_vec.shape[0] == 1
  
      mol_tree = DGLMolTree(None)
      mol_tree.graph = mol_tree.graph.to(mol_vec.device)
      mol_tree_graph = mol_tree.graph
  
      init_hidden = cuda(torch.zeros(1, self.hidden_size))
  
      root_hidden = torch.cat([init_hidden, mol_vec], 1)
      root_hidden = F.relu(self.W(root_hidden))
      root_score = self.W_o(root_hidden)
      _, root_wid = torch.max(root_score, 1)
      root_wid = root_wid.view(1)
  
      mol_tree_graph.add_nodes(1)  # root
      mol_tree_graph.ndata["wid"] = root_wid
      mol_tree_graph.ndata["x"] = self.embedding(root_wid)
      mol_tree_graph.ndata["h"] = init_hidden
      mol_tree_graph.ndata["fail"] = cuda(torch.tensor([0]))
      mol_tree.nodes_dict[0] = root_node_dict = create_node_dict(
          self.vocab.get_smiles(root_wid)
      )
  
      stack, trace = [], []
      stack.append((0, self.vocab.get_slots(root_wid)))
  
      all_nodes = {0: root_node_dict}
      h = {}
      first = True
      new_node_id = 0
      new_edge_id = 0
  
      for step in range(MAX_DECODE_LEN):
          u, u_slots = stack[-1]
          x = mol_tree_graph.ndata["x"][u : u + 1]
          h = mol_tree_graph.ndata["h"][u : u + 1]
  
          # Predict stop
          p_input = torch.cat([x, h, mol_vec], 1)
          p_score = torch.sigmoid(self.U_s(torch.relu(self.U(p_input))))
          backtrack = p_score.item() < 0.5
  
          if not backtrack:
              # Predict next clique.  Note that the prediction may fail due
              # to lack of assemblable components
              mol_tree_graph.add_nodes(1)
              new_node_id += 1
              v = new_node_id
              mol_tree_graph.add_edges(u, v)
              uv = new_edge_id
              new_edge_id += 1
  
              if first:
                  mol_tree_graph.edata.update(
                      {
                          "s": cuda(torch.zeros(1, self.hidden_size)),
                          "m": cuda(torch.zeros(1, self.hidden_size)),
                          "r": cuda(torch.zeros(1, self.hidden_size)),
                          "z": cuda(torch.zeros(1, self.hidden_size)),
                          "src_x": cuda(
                              torch.zeros(1, self.hidden_size)
                          ),
                          "dst_x": cuda(
                              torch.zeros(1, self.hidden_size)
                          ),
                          "rm": cuda(torch.zeros(1, self.hidden_size)),
                          "accum_rm": cuda(
                              torch.zeros(1, self.hidden_size)
                          ),
                      }
                  )
                  first = False
  
              mol_tree_graph.edata["src_x"][uv] = mol_tree_graph.ndata[
                  "x"
              ][u]
              # keeping dst_x 0 is fine as h on new edge doesn't depend on that.
  
              # DGL doesn't dynamically maintain a line graph.
              mol_tree_graph_lg = line_graph(
                  mol_tree_graph, backtracking=False, shared=True
              )
  
              mol_tree_graph_lg.pull(
                  uv, DGLF.copy_u("m", "m"), DGLF.sum("m", "s")
              )
              mol_tree_graph_lg.pull(
                  uv,
                  DGLF.copy_u("rm", "rm"),
                  DGLF.sum("rm", "accum_rm"),
              )
              mol_tree_graph_lg.apply_nodes(
                  self.dec_tree_edge_update.update_zm, v=uv
              )
              mol_tree_graph.edata.update(mol_tree_graph_lg.ndata)
              mol_tree_graph.pull(
                  v, DGLF.copy_e("m", "m"), DGLF.sum("m", "h")
              )
  
              h_v = mol_tree_graph.ndata["h"][v : v + 1]
              q_input = torch.cat([h_v, mol_vec], 1)
              q_score = torch.softmax(
                  self.W_o(torch.relu(self.W(q_input))), -1
              )
              _, sort_wid = torch.sort(q_score, 1, descending=True)
              sort_wid = sort_wid.squeeze()
  
              next_wid = None
              for wid in sort_wid.tolist()[:5]:
                  slots = self.vocab.get_slots(wid)
                  cand_node_dict = create_node_dict(
                      self.vocab.get_smiles(wid)
                  )
                  if have_slots(u_slots, slots) and can_assemble(
                      mol_tree, u, cand_node_dict
                  ):
                      next_wid = wid
                      next_slots = slots
                      next_node_dict = cand_node_dict
                      break
  
              if next_wid is None:
                  # Failed adding an actual children; v is a spurious node
                  # and we mark it.
                  mol_tree_graph.ndata["fail"][v] = cuda(
                      torch.tensor([1])
                  )
                  backtrack = True
              else:
                  next_wid = cuda(torch.tensor([next_wid]))
                  mol_tree_graph.ndata["wid"][v] = next_wid
                  mol_tree_graph.ndata["x"][v] = self.embedding(
                      next_wid
                  )
                  mol_tree.nodes_dict[v] = next_node_dict
                  all_nodes[v] = next_node_dict
                  stack.append((v, next_slots))
                  mol_tree_graph.add_edges(v, u)
                  vu = new_edge_id
                  new_edge_id += 1
                  mol_tree_graph.edata["dst_x"][
                      uv
                  ] = mol_tree_graph.ndata["x"][v]
                  mol_tree_graph.edata["src_x"][
                      vu
                  ] = mol_tree_graph.ndata["x"][v]
                  mol_tree_graph.edata["dst_x"][
                      vu
                  ] = mol_tree_graph.ndata["x"][u]
  
                  # DGL doesn't dynamically maintain a line graph.
                  mol_tree_graph_lg = line_graph(
                      mol_tree_graph, backtracking=False, shared=True
                  )
                  mol_tree_graph_lg.apply_nodes(
                      self.dec_tree_edge_update.update_r, uv
                  )
                  mol_tree_graph.edata.update(mol_tree_graph_lg.ndata)
  
          if backtrack:
              if len(stack) == 1:
                  break  # At root, terminate
  
              pu, _ = stack[-2]
              u_pu = mol_tree_graph.edge_ids(u, pu)
  
              mol_tree_graph_lg.pull(
                  u_pu, DGLF.copy_u("m", "m"), DGLF.sum("m", "s")
              )
              mol_tree_graph_lg.pull(
                  u_pu,
                  DGLF.copy_u("rm", "rm"),
                  DGLF.sum("rm", "accum_rm"),
              )
              mol_tree_graph_lg.apply_nodes(
                  self.dec_tree_edge_update, v=u_pu
              )
              mol_tree_graph.edata.update(mol_tree_graph_lg.ndata)
              mol_tree_graph.pull(
                  pu, DGLF.copy_e("m", "m"), DGLF.sum("m", "h")
              )
              stack.pop()
  
      effective_nodes = mol_tree_graph.filter_nodes(
          lambda nodes: nodes.data["fail"] != 1
      )
      effective_nodes, _ = torch.sort(effective_nodes)
      return mol_tree, all_nodes, effective_nodes
  
file_path: dgl/examples/pytorch/jtnn/jtnn/jtnn_dec.py
imports: import dgl.function
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import batch, dfs_labeled_edges_generator, line_graph
  from .chemutils import enum_assemble_nx, get_mol
  from .mol_tree_nx import DGLMolTree
  from .nnutils import cuda, GRUUpdate, tocpu
label_desc: Use image node_id 4 for calling the DGLJTNNDecoder obj's underlying member method code with example usage: obj.decode(mol_vec) and returns: mol_tree,  all_nodes,  effective_nodes