code: def embed_trees(self, inputs):
      if self.config["backprop"]:
          states = self.constituency_parser.analyze_trees(inputs)
      else:
          with torch.no_grad():
              states = self.constituency_parser.analyze_trees(inputs)
  
      constituent_lists = [x.constituents for x in states]
      states = [x.state for x in states]
  
      word_begin_hx = torch.stack(
          [state.word_queue[0].hx for state in states]
      )
      word_end_hx = torch.stack(
          [state.word_queue[state.word_position].hx for state in states]
      )
      transition_hx = torch.stack(
          [
              self.constituency_parser.transition_stack.output(
                  state.transitions
              )
              for state in states
          ]
      )
      # go down one layer to get the embedding off the top of the S, not the ROOT
      # (in terms of the typical treebank)
      # the idea being that the ROOT has no additional information
      # and may even have 0s for the embedding in certain circumstances,
      # such as after learning UNTIED_MAX long enough
      if self.config["top_layer"]:
          constituent_hx = torch.stack(
              [
                  self.constituency_parser.constituent_stack.output(
                      state.constituents
                  )
                  for state in states
              ]
          )
      else:
          constituent_hx = torch.cat(
              [
                  constituents[-2].tree_hx
                  for constituents in constituent_lists
              ],
              dim=0,
          )
  
      if self.config["all_words"]:
          # need B matrices of N x hidden_size
          key = [
              torch.stack(
                  [
                      torch.cat([word.hx, thx, chx])
                      for word in state.word_queue
                  ],
                  dim=0,
              )
              for state, thx, chx in zip(
                  states, transition_hx, constituent_hx
              )
          ]
      else:
          key = torch.cat(
              (
                  word_begin_hx,
                  word_end_hx,
                  transition_hx,
                  constituent_hx,
              ),
              dim=1,
          ).unsqueeze(1)
  
      if not self.config["node_attn"]:
          return key
      key = [self.key(x) for x in key]
  
      node_hx = [
          torch.stack([con.tree_hx for con in constituents], dim=0)
          for constituents in constituent_lists
      ]
      queries = [
          self.query(nhx).reshape(nhx.shape[0], -1) for nhx in node_hx
      ]
      values = [
          self.value(nhx).reshape(nhx.shape[0], -1) for nhx in node_hx
      ]
      # TODO: could pad to make faster here
      attn = [
          torch.matmul(q, k.transpose(0, 1))
          for q, k in zip(queries, key)
      ]
      attn = [torch.softmax(x, dim=0) for x in attn]
      previous_layer = [
          torch.matmul(weight.transpose(0, 1), value)
          for weight, value in zip(attn, values)
      ]
      return previous_layer
  
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 2 for calling the TreeEmbedding obj's underlying member method code with example usage: obj.embed_trees(inputs) and returns: previous_layer, key