code: def __init__(self, constituency_parser, args):
      super(TreeEmbedding, self).__init__()
  
      self.config = {
          "all_words": args["all_words"],
          "backprop": args["backprop"],
          # "batch_norm":  args["batch_norm"],
          "node_attn": args["node_attn"],
          "top_layer": args["top_layer"],
      }
  
      self.constituency_parser = constituency_parser
  
      # word_lstm:         hidden_size * num_tree_lstm_layers * 2 (start & end)
      # transition_stack:  transition_hidden_size
      # constituent_stack: hidden_size
      self.hidden_size = (
          self.constituency_parser.hidden_size
          + self.constituency_parser.transition_hidden_size
      )
      if self.config["all_words"]:
          self.hidden_size += (
              self.constituency_parser.hidden_size
              * self.constituency_parser.num_tree_lstm_layers
          )
      else:
          self.hidden_size += (
              self.constituency_parser.hidden_size
              * self.constituency_parser.num_tree_lstm_layers
              * 2
          )
  
      if self.config["node_attn"]:
          self.query = nn.Linear(
              self.constituency_parser.hidden_size,
              self.constituency_parser.hidden_size,
          )
          self.key = nn.Linear(
              self.hidden_size, self.constituency_parser.hidden_size
          )
          self.value = nn.Linear(
              self.constituency_parser.hidden_size,
              self.constituency_parser.hidden_size,
          )
  
          # TODO: cat transition and constituent hx as well?
          self.output_size = (
              self.constituency_parser.hidden_size
              * self.constituency_parser.num_tree_lstm_layers
          )
      else:
          self.output_size = self.hidden_size
  
file_path: stanza/stanza/models/constituency/tree_embedding.py
imports: import torch
  import torch.nn
  from stanza.models.constituency.trainer import Trainer
label_desc: Use image node_id 1 to create a new TreeEmbedding object from inherited base classes: nn with example: obj = TreeEmbedding(constituency_parser, args)