code: def run(self, mol_tree_batch, mol_tree_batch_lg):
      # Since tree roots are designated to 0.  In the batched graph we can
      # simply find the corresponding node ID by looking at node_offset
      node_offset = np.cumsum(
          np.insert(
              mol_tree_batch.batch_num_nodes().cpu().numpy(), 0, 0
          )
      )
      root_ids = node_offset[:-1]
      n_nodes = mol_tree_batch.num_nodes()
      n_edges = mol_tree_batch.num_edges()
  
      # Assign structure embeddings to tree nodes
      mol_tree_batch.ndata.update(
          {
              "x": self.embedding(mol_tree_batch.ndata["wid"]),
              "m": cuda(torch.zeros(n_nodes, self.hidden_size)),
              "h": cuda(torch.zeros(n_nodes, self.hidden_size)),
          }
      )
  
      # Initialize the intermediate variables according to Eq (4)-(8).
      # Also initialize the src_x and dst_x fields.
      # TODO: context?
      mol_tree_batch.edata.update(
          {
              "s": cuda(torch.zeros(n_edges, self.hidden_size)),
              "m": cuda(torch.zeros(n_edges, self.hidden_size)),
              "r": cuda(torch.zeros(n_edges, self.hidden_size)),
              "z": cuda(torch.zeros(n_edges, self.hidden_size)),
              "src_x": cuda(torch.zeros(n_edges, self.hidden_size)),
              "dst_x": cuda(torch.zeros(n_edges, self.hidden_size)),
              "rm": cuda(torch.zeros(n_edges, self.hidden_size)),
              "accum_rm": cuda(torch.zeros(n_edges, self.hidden_size)),
          }
      )
  
      # Send the source/destination node features to edges
      mol_tree_batch.apply_edges(
          func=lambda edges: {
              "src_x": edges.src["x"],
              "dst_x": edges.dst["x"],
          },
      )
  
      # Message passing
      # I exploited the fact that the reduce function is a sum of incoming
      # messages, and the uncomputed messages are zero vectors.  Essentially,
      # we can always compute s_ij as the sum of incoming m_ij, no matter
      # if m_ij is actually computed or not.
      mol_tree_batch_lg.ndata.update(mol_tree_batch.edata)
      for eid in level_order(mol_tree_batch, root_ids):
          eid = eid.to(mol_tree_batch_lg.device)
          mol_tree_batch_lg.pull(
              eid, DGLF.copy_u("m", "m"), DGLF.sum("m", "s")
          )
          mol_tree_batch_lg.pull(
              eid, DGLF.copy_u("rm", "rm"), DGLF.sum("rm", "accum_rm")
          )
          mol_tree_batch_lg.apply_nodes(self.enc_tree_update, v=eid)
  
      # Readout
      mol_tree_batch.edata.update(mol_tree_batch_lg.ndata)
      mol_tree_batch.update_all(
          DGLF.copy_e("m", "m"), DGLF.sum("m", "m")
      )
      mol_tree_batch.apply_nodes(self.enc_tree_gather_update)
  
      root_vecs = mol_tree_batch.nodes[root_ids].data["h"]
  
      return mol_tree_batch, root_vecs
  
file_path: dgl/examples/pytorch/jtnn/jtnn/jtnn_enc.py
imports: import dgl.function
  import numpy
  import torch
  import torch.nn
  from dgl import batch, bfs_edges_generator, line_graph
  from .nnutils import cuda, GRUUpdate, tocpu
label_desc: Use image node_id 3 for calling the DGLJTNNEncoder obj's underlying member method code with example usage: obj.run(mol_tree_batch, mol_tree_batch_lg) and returns: mol_tree_batch,  root_vecs