code: def test_push_state_same_length():
      """
      Test that stacks are being updated correctly when using 3 stacks of the same length
  
      Values of the attention are not verified, though
      """
      ts = TransformerTreeStack(3, 5, 0.0)
      initial = ts.initial_state()
      rand_input = torch.randn(3, 3)
      stacks = ts.push_states(
          [initial, initial, initial], ["A", "A", "A"], rand_input
      )
      stacks = ts.push_states(stacks, ["B", "B", "B"], rand_input)
      stacks = ts.push_states(stacks, ["C", "C", "C"], rand_input)
      assert len(stacks) == 3
      for s in stacks:
          assert len(s) == 4
          assert s.value.key_stack.shape == torch.Size([4, 5])
          assert s.value.value_stack.shape == torch.Size([4, 5])
          assert s.value.value == "C"
          assert s.pop().value.value == "B"
          assert s.pop().pop().value.value == "A"
          assert s.pop().pop().pop().value.value is None
  
file_path: stanza/stanza/tests/constituency/test_transformer_tree_stack.py
imports: import pytest
  import torch
  from stanza.models.constituency.transformer_tree_stack import TransformerTreeStack
label_desc: Use image node_id 4 for calling a global function with example usage: test_push_state_same_length() without return types