code: def test_mask():
      """
      Test that a mask prevents the softmax from picking up unwanted values
      """
      ts = TransformerTreeStack(3, 5, 0.0)
  
      random_v = torch.tensor([[[0.1, 0.2, 0.3, 0.4, 0.5]]])
      double_v = random_v * 2
      value = torch.cat([random_v, double_v], axis=1)
      random_k = torch.randn(1, 1, 5)
      key = torch.cat([random_k, random_k], axis=1)
      query = torch.randn(1, 5)
  
      output = ts.attention(key, query, value)
      # when the two keys are equal, we expect the attention to be 50/50
      expected_output = (random_v + double_v) / 2
      assert torch.allclose(output, expected_output)
  
      # If the first entry is masked out, the second one should be the
      # only one represented
      mask = torch.zeros(1, 2, dtype=torch.bool)
      mask[0][0] = True
      output = ts.attention(key, query, value, mask)
      assert torch.allclose(output, double_v)
  
      # If the second entry is masked out, the first one should be the
      # only one represented
      mask = torch.zeros(1, 2, dtype=torch.bool)
      mask[0][1] = True
      output = ts.attention(key, query, value, mask)
      assert torch.allclose(output, random_v)
  
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 6 for calling a global function with example usage: test_mask() without return types