docstr: Push new inputs to the stacks and rerun attention on them
  
  Where B is the number of items stacked together, I is input_size
  stacks: B TreeStacks such as produced by initial_state and/or push_states
  values: the new items to push on the stacks such as tree nodes or anything
  inputs: BxI for the new input items
  
  Runs attention starting from the existing keys & values
code: def push_states(self, stacks, values, inputs):
      """
      Push new inputs to the stacks and rerun attention on them
  
      Where B is the number of items stacked together, I is input_size
      stacks: B TreeStacks such as produced by initial_state and/or push_states
      values: the new items to push on the stacks such as tree nodes or anything
      inputs: BxI for the new input items
  
      Runs attention starting from the existing keys & values
      """
      device = self.w_key.weight.device
  
      batch_len = len(stacks)  # B
      positions = [x.value.key_stack.shape[0] for x in stacks]
      max_len = max(positions)  # N
  
      if self.use_position:
          position_encodings = self.position_encoding(positions)
          inputs = inputs + position_encodings
  
      inputs = self.input_dropout(inputs)
      if len(inputs.shape) == 3:
          if inputs.shape[0] == 1:
              inputs = inputs.squeeze(0)
          else:
              raise ValueError(
                  "Expected the inputs to be of shape 1xBxI, got {}".format(
                      inputs.shape
                  )
              )
  
      new_keys = self.w_key(inputs)
      key_stack = torch.zeros(
          batch_len, max_len + 1, self.output_size, device=device
      )
      key_stack[:, -1, :] = new_keys
      for stack_idx, stack in enumerate(stacks):
          key_stack[
              stack_idx, -positions[stack_idx] - 1 : -1, :
          ] = stack.value.key_stack
  
      new_values = self.w_value(inputs)
      value_stack = torch.zeros(
          batch_len, max_len + 1, self.output_size, device=device
      )
      value_stack[:, -1, :] = new_values
      for stack_idx, stack in enumerate(stacks):
          value_stack[
              stack_idx, -positions[stack_idx] - 1 : -1, :
          ] = stack.value.value_stack
  
      query = self.w_query(inputs)
  
      mask = torch.zeros(
          batch_len, max_len + 1, device=device, dtype=torch.bool
      )
      for stack_idx, stack in enumerate(stacks):
          if len(stack) < max_len:
              masked = max_len - positions[stack_idx]
              mask[stack_idx, :masked] = True
  
      batched_output = self.attention(
          key_stack, query, value_stack, mask
      )
  
      new_stacks = []
      for stack_idx, (
          stack,
          node_value,
          new_key,
          new_value,
          output,
      ) in enumerate(
          zip(stacks, values, key_stack, value_stack, batched_output)
      ):
          # max_len-len(stack) so that we ignore the padding at the start of shorter stacks
          new_key_stack = new_key[max_len - positions[stack_idx] :, :]
          new_value_stack = new_value[
              max_len - positions[stack_idx] :, :
          ]
          if (
              self.length_limit is not None
              and new_key_stack.shape[0] > self.length_limit + 1
          ):
              new_key_stack = torch.cat(
                  [new_key_stack[:1, :], new_key_stack[2:, :]], axis=0
              )
              new_value_stack = torch.cat(
                  [new_value_stack[:1, :], new_value_stack[2:, :]],
                  axis=0,
              )
          new_stacks.append(
              stack.push(
                  value=Node(
                      node_value, new_key_stack, new_value_stack, output
                  )
              )
          )
      return new_stacks
  
file_path: stanza/stanza/models/constituency/transformer_tree_stack.py
imports: from collections import namedtuple
  import torch
  import torch.nn
  from stanza.models.constituency.positional_encoding import SinusoidalEncoding
  from stanza.models.constituency.tree_stack import TreeStack
label_desc: Use image node_id 4 for calling the TransformerTreeStack obj's underlying member method code with example usage: obj.push_states(stacks, values, inputs) and returns: new_stacks