docstr: Build new constituents with the given label from the list of children
  
  labels is a list of labels for each of the new nodes to construct
  children_lists is a list of children that go under each of the new nodes
  lists of each are used so that we can stack operations
code: def build_constituents(self, labels, children_lists):
      """
      Build new constituents with the given label from the list of children
  
      labels is a list of labels for each of the new nodes to construct
      children_lists is a list of children that go under each of the new nodes
      lists of each are used so that we can stack operations
      """
      # at the end of each of these operations, we expect lstm_hx.shape
      # is (L, N, hidden_size) for N lists of children
      if (
          self.constituency_composition
          == ConstituencyComposition.BILSTM
          or self.constituency_composition
          == ConstituencyComposition.BILSTM_MAX
      ):
          node_hx = [
              [child.value.tree_hx.squeeze(0) for child in children]
              for children in children_lists
          ]
          label_hx = [
              self.constituent_open_embedding(
                  self.constituent_open_tensors[
                      self.constituent_open_map[label]
                  ]
              )
              for label in labels
          ]
  
          max_length = max(len(children) for children in children_lists)
          zeros = torch.zeros(
              self.hidden_size, device=label_hx[0].device
          )
          # weirdly, this is faster than using pack_sequence
          unpacked_hx = [
              [lhx] + nhx + [lhx] + [zeros] * (max_length - len(nhx))
              for lhx, nhx in zip(label_hx, node_hx)
          ]
          unpacked_hx = [
              self.lstm_input_dropout(torch.stack(nhx))
              for nhx in unpacked_hx
          ]
          packed_hx = torch.stack(unpacked_hx, axis=1)
          packed_hx = torch.nn.utils.rnn.pack_padded_sequence(
              packed_hx,
              [len(x) + 2 for x in children_lists],
              enforce_sorted=False,
          )
          lstm_output = self.constituent_reduce_lstm(packed_hx)
          # take just the output of the final layer
          #   result of lstm is ouput, (hx, cx)
          #   so [1][0] gets hx
          #      [1][0][-1] is the final output
          # will be shape len(children_lists) * 2, hidden_size for bidirectional
          # where forward outputs are -2 and backwards are -1
          if (
              self.constituency_composition
              == ConstituencyComposition.BILSTM
          ):
              lstm_output = lstm_output[1][0]
              forward_hx = lstm_output[-2, :, :]
              backward_hx = lstm_output[-1, :, :]
              hx = self.reduce_linear(
                  torch.cat((forward_hx, backward_hx), axis=1)
              )
          else:
              (
                  lstm_output,
                  lstm_lengths,
              ) = torch.nn.utils.rnn.pad_packed_sequence(lstm_output[0])
              lstm_output = [
                  lstm_output[1 : length - 1, x, :]
                  for x, length in zip(
                      range(len(lstm_lengths)), lstm_lengths
                  )
              ]
              lstm_output = torch.stack(
                  [torch.max(x, 0).values for x in lstm_output], axis=0
              )
              hx = self.reduce_forward(
                  lstm_output[:, : self.hidden_size]
              ) + self.reduce_backward(
                  lstm_output[:, self.hidden_size :]
              )
          lstm_hx = self.nonlinearity(hx).unsqueeze(0)
          lstm_cx = None
      elif self.constituency_composition == ConstituencyComposition.MAX:
          node_hx = [
              [child.value.tree_hx for child in children]
              for children in children_lists
          ]
          unpacked_hx = [
              self.lstm_input_dropout(
                  torch.max(torch.stack(nhx), 0).values
              )
              for nhx in node_hx
          ]
          packed_hx = torch.stack(unpacked_hx, axis=1)
          hx = self.reduce_linear(packed_hx)
          lstm_hx = self.nonlinearity(hx)
          lstm_cx = None
      elif (
          self.constituency_composition
          == ConstituencyComposition.UNTIED_MAX
      ):
          node_hx = [
              [child.value.tree_hx for child in children]
              for children in children_lists
          ]
          unpacked_hx = [
              self.lstm_input_dropout(
                  torch.max(torch.stack(nhx), 0).values
              )
              for nhx in node_hx
          ]
          # shape == len(labels),1,hidden_size after the stack
          # packed_hx = torch.stack(unpacked_hx, axis=0)
          label_indices = [
              self.constituent_open_map[label] for label in labels
          ]
          # we would like to stack the reduce_linear_weight calculations as follows:
          # reduce_weight = self.reduce_linear_weight[label_indices]
          # reduce_bias = self.reduce_linear_bias[label_indices]
          # this would allow for faster vectorized operations.
          # however, this runs out of memory on larger training examples,
          # presumably because there are too many stacks in a row and each one
          # has its own gradient kept for the entire calculation
          # fortunately, this operation is not a huge part of the expense
          hx = [
              torch.matmul(
                  self.reduce_linear_weight[label_idx],
                  hx_layer.squeeze(0),
              )
              + self.reduce_linear_bias[label_idx]
              for label_idx, hx_layer in zip(label_indices, unpacked_hx)
          ]
          hx = torch.stack(hx, axis=0)
          hx = hx.unsqueeze(0)
          lstm_hx = self.nonlinearity(hx)
          lstm_cx = None
      elif (
          self.constituency_composition
          == ConstituencyComposition.BIGRAM
      ):
          node_hx = [
              [child.value.tree_hx for child in children]
              for children in children_lists
          ]
          unpacked_hx = []
          for nhx in node_hx:
              # tanh or otherwise limit the size of the output?
              stacked_nhx = self.lstm_input_dropout(
                  torch.cat(nhx, axis=0)
              )
              if stacked_nhx.shape[0] > 1:
                  bigram_hx = torch.cat(
                      (stacked_nhx[:-1, :], stacked_nhx[1:, :]), axis=1
                  )
                  bigram_hx = self.reduce_bigram(bigram_hx) / 2
                  stacked_nhx = torch.cat(
                      (stacked_nhx, bigram_hx), axis=0
                  )
              unpacked_hx.append(torch.max(stacked_nhx, 0).values)
          packed_hx = torch.stack(unpacked_hx, axis=0).unsqueeze(0)
          hx = self.reduce_linear(packed_hx)
          lstm_hx = self.nonlinearity(hx)
          lstm_cx = None
      elif (
          self.constituency_composition == ConstituencyComposition.ATTN
      ):
          node_hx = [
              [child.value.tree_hx for child in children]
              for children in children_lists
          ]
          label_hx = [
              self.constituent_open_embedding(
                  self.constituent_open_tensors[
                      self.constituent_open_map[label]
                  ]
              )
              for label in labels
          ]
          unpacked_hx = [torch.stack(nhx) for nhx in node_hx]
          unpacked_hx = [
              torch.cat((lhx.unsqueeze(0).unsqueeze(0), nhx), axis=0)
              for lhx, nhx in zip(label_hx, unpacked_hx)
          ]
          unpacked_hx = [
              self.reduce_attn(nhx, nhx, nhx)[0].squeeze(1)
              for nhx in unpacked_hx
          ]
          unpacked_hx = [
              self.lstm_input_dropout(torch.max(nhx, 0).values)
              for nhx in unpacked_hx
          ]
          hx = torch.stack(unpacked_hx, axis=0)
          lstm_hx = self.nonlinearity(hx).unsqueeze(0)
          lstm_cx = None
      elif (
          self.constituency_composition == ConstituencyComposition.KEY
          or self.constituency_composition
          == ConstituencyComposition.UNTIED_KEY
      ):
          node_hx = [
              torch.stack([child.value.tree_hx for child in children])
              for children in children_lists
          ]
          # add a position vector to each node_hx
          node_hx = [
              self.reduce_position(x.reshape(x.shape[0], -1))
              for x in node_hx
          ]
          query_hx = [self.reduce_query(nhx) for nhx in node_hx]
          # reshape query for MHA
          query_hx = [
              nhx.reshape(
                  nhx.shape[0], self.reduce_heads, -1
              ).transpose(0, 1)
              for nhx in query_hx
          ]
          if (
              self.constituency_composition
              == ConstituencyComposition.KEY
          ):
              queries = [
                  torch.matmul(nhx, self.reduce_key) for nhx in query_hx
              ]
          else:
              label_indices = [
                  self.constituent_open_map[label] for label in labels
              ]
              queries = [
                  torch.matmul(nhx, self.reduce_key[label_idx])
                  for nhx, label_idx in zip(query_hx, label_indices)
              ]
          # softmax each head
          weights = [
              torch.nn.functional.softmax(nhx, dim=1).transpose(1, 2)
              for nhx in queries
          ]
          value_hx = [self.reduce_value(nhx) for nhx in node_hx]
          value_hx = [
              nhx.reshape(
                  nhx.shape[0], self.reduce_heads, -1
              ).transpose(0, 1)
              for nhx in value_hx
          ]
          # use the softmaxes to add up the heads
          unpacked_hx = [
              torch.matmul(weight, nhx).squeeze(1)
              for weight, nhx in zip(weights, value_hx)
          ]
          unpacked_hx = [nhx.reshape(-1) for nhx in unpacked_hx]
          hx = torch.stack(unpacked_hx, axis=0).unsqueeze(0)
          lstm_hx = self.nonlinearity(hx)
          lstm_cx = None
      elif self.constituency_composition in (
          ConstituencyComposition.TREE_LSTM,
          ConstituencyComposition.TREE_LSTM_CX,
      ):
          label_hx = [
              self.lstm_input_dropout(
                  self.constituent_open_embedding(
                      self.constituent_open_tensors[
                          self.constituent_open_map[label]
                      ]
                  )
              )
              for label in labels
          ]
          label_hx = torch.stack(label_hx).unsqueeze(0)
  
          max_length = max(len(children) for children in children_lists)
  
          # stacking will let us do elementwise multiplication faster, hopefully
          node_hx = [
              [child.value.tree_hx for child in children]
              for children in children_lists
          ]
          unpacked_hx = [
              self.lstm_input_dropout(torch.stack(nhx))
              for nhx in node_hx
          ]
          unpacked_hx = [nhx.max(dim=0) for nhx in unpacked_hx]
          packed_hx = torch.stack(
              [nhx.values for nhx in unpacked_hx], axis=1
          )
          # packed_hx = packed_hx.max(dim=0).values
  
          node_cx = [
              torch.stack([child.value.tree_cx for child in children])
              for children in children_lists
          ]
          node_cx_indices = [
              uhx.indices.unsqueeze(0) for uhx in unpacked_hx
          ]
          unpacked_cx = [
              ncx.gather(0, nci).squeeze(0)
              for ncx, nci in zip(node_cx, node_cx_indices)
          ]
          packed_cx = torch.stack(unpacked_cx, axis=1)
  
          _, (lstm_hx, lstm_cx) = self.constituent_reduce_lstm(
              label_hx, (packed_hx, packed_cx)
          )
      else:
          raise ValueError(
              "Unhandled ConstituencyComposition: {}".format(
                  self.constituency_composition
              )
          )
  
      constituents = []
      for idx, (label, children) in enumerate(
          zip(labels, children_lists)
      ):
          children = [child.value.value for child in children]
          if isinstance(label, str):
              node = Tree(label=label, children=children)
          else:
              for value in reversed(label):
                  node = Tree(label=value, children=children)
                  children = node
          constituents.append(
              Constituent(
                  node,
                  lstm_hx[:, idx, :],
                  lstm_cx[:, idx, :] if lstm_cx is not None else None,
              )
          )
      return constituents
  
imports: from collections import namedtuple
  from enum import Enum
  import logging
  import math
  import random
  import torch
  import torch.nn
  from torch.nn.utils.rnn import pack_padded_sequence
  from stanza.models.common.bert_embedding import extract_bert_embeddings
  from stanza.models.common.maxout_linear import MaxoutLinear
  from stanza.models.common.utils import unsort
  from stanza.models.common.vocab import PAD_ID, UNK_ID
  from stanza.models.constituency.base_model import BaseModel
  from stanza.models.constituency.label_attention import LabelAttentionModule
  from stanza.models.constituency.lstm_tree_stack import LSTMTreeStack
  from stanza.models.constituency.parse_transitions import TransitionScheme
  from stanza.models.constituency.parse_tree import Tree
  from stanza.models.constituency.partitioned_transformer import PartitionedTransformerModule
  from stanza.models.constituency.positional_encoding import ConcatSinusoidalEncoding
  from stanza.models.constituency.transformer_tree_stack import TransformerTreeStack
  from stanza.models.constituency.tree_stack import TreeStack
  from stanza.models.constituency.utils import build_nonlinearity, initialize_linear
label_desc: Use image node_id 21 for calling the LSTMModel obj's underlying member method code with example usage: obj.build_constituents(labels, children_lists) and returns: constituents