docstr: Compute tree-lstm prediction given a batch.
  
  Parameters
  ----------
  batch : dgl.data.SSTBatch
      The data batch.
  h : Tensor
      Initial hidden state.
  c : Tensor
      Initial cell state.
  
  Returns
  -------
  logits : Tensor
      The prediction of each node.
docstr_tok: ["Compute","tree-lstm","prediction","given","a","batch",".","Parameters","--","--","--","--","--","batch",":","dgl.data.SSTBatch","The","data","batch",".","h",":","Tensor","Initial","hidden","state",".","c",":","Tensor","Initial","cell","state",".","Returns","--","--","--","-","logits",":","Tensor","The","prediction","of","each","node","."]
code: def forward(self, batch, h, c):
      """Compute tree-lstm prediction given a batch.
  
      Parameters
      ----------
      batch : dgl.data.SSTBatch
          The data batch.
      h : Tensor
          Initial hidden state.
      c : Tensor
          Initial cell state.
  
      Returns
      -------
      logits : Tensor
          The prediction of each node.
      """
      g = batch.graph
      # to heterogenous graph
      g = dgl.graph(g.edges())
      # feed embedding
      embeds = self.embedding(batch.wordid * batch.mask)
      g.ndata["iou"] = self.cell.W_iou(
          self.dropout(embeds)
      ) * batch.mask.float().unsqueeze(-1)
      g.ndata["h"] = h
      g.ndata["c"] = c
      # propagate
      dgl.prop_nodes_topo(
          g,
          message_func=self.cell.message_func,
          reduce_func=self.cell.reduce_func,
          apply_node_func=self.cell.apply_node_func,
      )
      # compute logits
      h = self.dropout(g.ndata.pop("h"))
      logits = self.linear(h)
      return logits
  
imports: import os
  from collections import namedtuple
  import dgl
  from dgl.data.tree import SSTDataset
  import matplotlib.pyplot
  import networkx
  import torch
  import torch.nn
  import dgl.function
  import torch
  import torch.nn.functional
  from torch.utils.data import DataLoader
label_desc: Use image node_id 2 for calling the TreeLSTM obj's underlying member method code with example usage: obj.forward(batch, h, c) and returns: logits