Attention - version 23#

This page documents version 23 of operator Attention. See Attention for the latest version (since version 24).

  • Domain: ai.onnx

  • Since version: 23

Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed.

This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V.

For self attention, kv_sequence_length equals to q_sequence_length.

For cross attention, query and key might have different lengths.

This operator also covers the 3 following variants based on the number of heads:

  1. Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, q_num_heads = kv_num_heads.

  2. Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, q_num_heads > kv_num_heads, q_num_heads % kv_num_heads == 0.

  3. Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, q_num_heads > kv_num_heads, kv_num_heads=1.

Attention bias to be added is calculated based on attn_mask input and is_causal attribute:

  1. If is_causal is set to 1, a query index i attends keys j <= i + past_sequence_length.

  2. attn_mask: A boolean mask where a value of True indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.

  3. If both attn_mask and is_causal are set, the valid positions are the intersection of both masks.

If a query row is fully masked after this intersection, its output row is zero.

Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided:

  The following pattern is applied by this operator:
      Q          K          V
      |          |          |
Q*sqrt(scale) K*sqrt(scale) |
      |          |          |
      |       Transpose     |
      |          |          |
      ---MatMul---          |
            |               |
  softcap (if provided)     |
            |               |
 at_mask---Add              |
            |               |
         Softmax            |
            |               |
            -----MatMul------
                   |
                   Y

Inputs

  • Q (T1): Query tensor. 4D tensor with shape (batch_size, q_num_heads, q_sequence_length, head_size) or 3D tensor with shape (batch_size, q_sequence_length, q_hidden_size). For cases with a 3D input tensor, q_hidden_size = q_num_heads * head_size

  • K (T1): Key tensor. 4D tensor with shape (batch_size, kv_num_heads, kv_sequence_length, head_size) or 3D tensor with shape (batch_size, kv_sequence_length, k_hidden_size). For cases with a 3D input tensor, k_hidden_size = kv_num_heads * head_size

  • V (T2): Value tensor. 4D tensor with shape (batch_size, kv_num_heads, kv_sequence_length, v_head_size) or 3D tensor with shape (batch_size, kv_sequence_length, v_hidden_size). For cases with a 3D input tensor, v_hidden_size = kv_num_heads * v_head_size

  • attn_mask (U): Attention mask. Shape must be broadcastable to 4D tensor with shape (batch_size, q_num_heads, q_sequence_length, total_sequence_length) where total_sequence_length = past_sequence_length + kv_sequence_length. Two types of masks are supported. A boolean mask where a value of True indicates that the element should take part in attention. Also supports a float mask of the same type as query, key, value that is added to the attention score.

  • past_key (T1): past state cache for key with shape (batch_size, kv_num_heads, past_sequence_length, head_size)

  • past_value (T2): past state cache for value with shape (batch_size, kv_num_heads, past_sequence_length, v_head_size)

Outputs

  • Y (T1): The output tensor . 4D tensor with shape (batch_size, q_num_heads, q_sequence_length, v_head_size) or 3D tensor with shape (batch_size, q_sequence_length, hidden_size). For cases with a 3D input tensor, hidden_size = q_num_heads * v_head_size

  • present_key (T1): Updated key cache with shape (batch_size, kv_num_heads, total_sequence_length, head_size) where total_sequence_length = past_sequence_length + kv_sequence_length.

  • present_value (T2): Updated value cache with shape (batch_size, kv_num_heads, total_sequence_length, v_head_size) where total_sequence_length = past_sequence_length + kv_sequence_length.

  • qk_matmul_output (T1): The output of QK matmul. 4D tensor with shape (batch_size, q_num_heads, q_sequence_length, total_sequence_length) where total_sequence_length = past_sequence_length + kv_sequence_length.

Type Constraints

  • T1: Constrain Q and K inputs types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).

  • T2: Constrain V input types to float tensors. Allowed types: tensor(bfloat16), tensor(double), tensor(float), tensor(float16).

  • U: Constrain output ‘mask’ types to boolean tensors and input types. Allowed types: tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8).

Examples#

test_cc_attention_23_boolmask_fullymasked_row_nan_robustness

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=bool
    [[[[False, False, False],
       [ True,  True,  True]],

      [[False, False, False],
       [ True,  True,  True]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.        ,  0.        ],
       [-0.16495374,  0.7457248 ]],

      [[ 0.        ,  0.        ],
       [ 0.18005186, -0.17207974]]]]

test_cc_attention_23_fullymasked_qk_matmul_output_mode3_zero

Node:
  Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output)
  Attributes:
    qk_matmul_output_mode = 3
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=bool
    [[[[False, False, False],
       [False, False, False]],

      [[False, False, False],
       [False, False, False]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[0., 0.],
       [0., 0.]],

      [[0., 0.],
       [0., 0.]]]]
  qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
    [[[[nan, nan, nan],
       [nan, nan, nan]],

      [[nan, nan, nan],
       [nan, nan, nan]]]]

test_cc_attention_3d

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.23092115,  0.5444725 ,  0.6482418 , -0.37858847],
      [-0.23092115,  0.77539366,  0.04638294, -0.08028026]]]

test_cc_attention_3d_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.458196  ,  0.41020232,  0.9921614 , -0.7460807 ],
      [ 0.0697852 ,  0.6150191 ,  0.3994022 , -0.34422377]]]

test_cc_attention_3d_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    is_causal = 1
Inputs:
  Q: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.  ,  1.  ,  1.  , -1.  ],
      [ 0.5 ,  0.5 ,  0.25,  0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 3, 4), dtype=float32
    [[[ 1.        ,  0.        ,  2.        , -2.        ],
      [ 0.412521  ,  0.587479  ,  0.79335546, -0.19003323],
      [ 0.        ,  0.6666667 ,  0.7425821 , -0.5205673 ]]]

test_cc_attention_3d_diff_heads_sizes

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]

Outputs:
  Y: shape=(1, 2, 6), dtype=float32
    [[[ 0.23092115,  0.5444725 ,  0.296508  ,  0.6482418 , -0.37858847,
        0.3571369 ],
      [-0.23092115,  0.77539366,  0.64288974,  0.04638294, -0.08028026,
        0.57520705]]]

test_cc_attention_3d_diff_heads_sizes_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 6), dtype=float32
    [[[ 0.458196  ,  0.41020232,  0.03320454,  0.9921614 , -0.7460807 ,
        0.4670852 ],
      [ 0.0697852 ,  0.6150191 ,  0.3722638 ,  0.3994022 , -0.34422377,
        0.5532167 ]]]

test_cc_attention_3d_diff_heads_sizes_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    is_causal = 1
Inputs:
  Q: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.  ,  1.  ,  1.  , -1.  ],
      [ 0.5 ,  0.5 ,  0.25,  0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]

Outputs:
  Y: shape=(1, 3, 6), dtype=float32
    [[[ 1.        ,  0.        , -1.        ,  2.        , -2.        ,
        1.        ],
      [ 0.412521  ,  0.587479  ,  0.762437  ,  0.79335546, -0.19003323,
       -0.0055371 ],
      [ 0.        ,  0.6666667 ,  0.5       ,  0.7425821 , -0.5205673 ,
        0.43631145]]]

test_cc_attention_3d_diff_heads_sizes_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]

Outputs:
  Y: shape=(1, 2, 6), dtype=float32
    [[[ 0.00333332,  0.6649986 ,  0.49749377,  0.6665943 , -0.58079195,
        0.5803768 ],
      [-0.00333332,  0.6683319 ,  0.50249374,  0.65489024, -0.5724745 ,
        0.5816217 ]]]

test_cc_attention_3d_diff_heads_sizes_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]

Outputs:
  Y: shape=(1, 2, 6), dtype=float32
    [[[ 0.15169363,  0.6033754 ,  0.44272968,  0.67558366, -0.4552321 ,
        0.4223395 ],
      [-0.15169363,  0.755069  ,  0.67027014,  0.29168454, -0.3030643 ,
        0.60680556]]]

test_cc_attention_3d_diff_heads_with_past_and_present

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 6), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  2.  , -2.  ,  1.  ],
      [ 0.  ,  1.  ,  2.  ,  0.5 ,  0.25, -0.25],
      [-1.  ,  1.  ,  0.5 , -0.5 ,  0.  ,  1.  ]]]
  attn_mask: shape=(2, 5), dtype=float32
    [[ 0. , -0.5, -1. ,  0.2,  0. ],
     [ 0.5,  0. , -0.2, -0.1,  0. ]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.5 ,  0.5 , -1.  ],
       [ 0.  ,  0.25,  0.5 ]],

      [[ 0.  ,  0.5 ,  0.5 ],
       [-0.5 ,  0.75, -0.25]]]]

Outputs:
  Y: shape=(1, 2, 6), dtype=float32
    [[[ 0.04906689,  0.6637103 ,  0.40457496,  0.17558974,  0.17481029,
        0.24073282],
      [-0.03033203,  0.6089026 ,  0.27225977,  0.01557301,  0.22852218,
        0.47354427]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 3), dtype=float32
    [[[[ 0.5 ,  0.5 , -1.  ],
       [ 0.  ,  0.25,  0.5 ],
       [ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 0.  ,  0.5 ,  0.5 ],
       [-0.5 ,  0.75, -0.25],
       [ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]

test_cc_attention_3d_gqa

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 8), dtype=float32
    [[[-0.02356532,  0.6783799 , -0.03533878,  0.6841799 ,  0.6482418 ,
       -0.37858847,  0.37784207, -0.12898168],
      [-0.02356531,  0.6783799 ,  0.11724145,  0.6063233 ,  0.9917567 ,
       -0.74587834,  0.29831943, -0.26321504]]]

test_cc_attention_3d_gqa_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 8), dtype=float32
    [[[ 0.30531266,  0.50214726,  0.29782698,  0.50647473,  0.9921614 ,
       -0.7460807 ,  0.74361753, -0.47596362],
      [ 0.2229336 ,  0.5335671 ,  0.31907293,  0.4797093 ,  1.1873223 ,
       -1.0039468 ,  0.6715176 , -0.5857588 ]]]

test_cc_attention_3d_gqa_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
    is_causal = 1
Inputs:
  Q: shape=(1, 3, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ],
      [ 0.2 , -0.1 ,  0.25,  0.  ,  0.5 ,  0.5 , -1.  ,  1.  ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 3, 8), dtype=float32
    [[[ 1.        ,  0.        ,  1.        ,  0.        ,  2.        ,
       -2.        ,  2.        , -2.        ],
      [ 0.4911621 ,  0.50883794,  0.5440794 ,  0.45592055,  1.25      ,
       -0.875     ,  0.9953577 , -0.4930365 ],
      [ 0.07057842,  0.63075423,  0.05884897,  0.636809  ,  0.6482418 ,
       -0.37858847,  1.478919  , -1.3989784 ]]]

test_cc_attention_3d_gqa_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 2, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 8), dtype=float32
    [[[-3.3333333e-04,  6.6683334e-01, -4.9999997e-04,  6.6691661e-01,
        6.6659433e-01, -5.8079195e-01,  6.6071773e-01, -5.7536310e-01],
      [-3.3333330e-04,  6.6683334e-01,  1.6666650e-03,  6.6583300e-01,
        6.7248535e-01, -5.8624268e-01,  6.6077113e-01, -5.7789290e-01]]]

test_cc_attention_3d_gqa_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 2, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 8), dtype=float32
    [[[-0.03040495,  0.6819007 , -0.04947086,  0.69103485,  0.67558366,
       -0.4552321 ,  0.40459138, -0.23742941],
      [-0.02114749,  0.67739695,  0.15203254,  0.5866078 ,  0.94510204,
       -0.72255105,  0.35040644, -0.32931072]]]

test_cc_attention_3d_gqa_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 8), dtype=float32
    [[[ 0.1 ,  0.2 , -0.1 ,  0.05,  0.5 ,  0.5 ,  1.  ,  0.  ],
      [ 0.3 ,  0.4 ,  0.2 , -0.3 ,  0.  ,  1.  ,  0.5 , -0.5 ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 8), dtype=float32
    [[[-0.12450287,  0.51090705, -0.13178737,  0.5070071 ,  0.48057237,
       -0.20996696,  0.27224937,  0.02852735],
      [-0.12746866,  0.51840895,  0.00804363,  0.47238812,  0.7363033 ,
       -0.5218315 ,  0.23465633, -0.06618155]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_3d_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.00333332,  0.6649986 ,  0.6665943 , -0.58079195],
      [-0.00333332,  0.6683319 ,  0.65489024, -0.5724745 ]]]

test_cc_attention_3d_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.15169363,  0.6033754 ,  0.67558366, -0.4552321 ],
      [-0.15169363,  0.755069  ,  0.29168454, -0.3030643 ]]]

test_cc_attention_3d_transpose_verification

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    q_num_heads = 3
    kv_num_heads = 3
Inputs:
  Q: shape=(1, 2, 12), dtype=float32
    [[[1., 1., 1., 1., 2., 2., 2., 2., 3., 3., 3., 3.],
      [1., 1., 1., 1., 2., 2., 2., 2., 3., 3., 3., 3.]]]
  K: shape=(1, 2, 12), dtype=float32
    [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
      [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]]
  V: shape=(1, 2, 12), dtype=float32
    [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
      [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]]

Outputs:
  Y: shape=(1, 2, 12), dtype=float32
    [[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
      [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]]]

test_cc_attention_3d_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.10764056,  0.45608166,  0.48057237, -0.20996696],
      [-0.31939295,  0.57814157,  0.05822894,  0.11373253]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_3d_with_past_and_present_qk_matmul

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.10764056,  0.45608166,  0.48057237, -0.20996696],
      [-0.31939295,  0.57814157,  0.05822894,  0.11373253]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.35355338,  0.        ,  0.70710677,  0.35355338,  0.        ],
       [-0.35355338,  0.35355338,  0.        ,  0.35355338,  0.70710677]],

      [[ 0.35355338,  0.17677669,  0.        ,  0.70710677, -0.08838835],
       [ 0.70710677, -1.0606601 , -1.4142135 ,  0.        ,  0.53033006]]]]

test_cc_attention_3d_with_past_and_present_qk_matmul_bias

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 0.5
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  attn_mask: shape=(2, 5), dtype=float32
    [[ 0. , -0.5, -1. ,  0.2,  0. ],
     [ 0.5,  0. , -0.2, -0.1,  0. ]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[-0.07067307,  0.63377535,  0.30545416,  0.01247976],
      [-0.23752189,  0.5571051 ,  0.09629637,  0.127618  ]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_3d_with_past_and_present_qk_matmul_softcap

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    scale = 1.0
    softcap = 0.5
    qk_matmul_output_mode = 2
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.05362739,  0.4880938 ,  0.4812978 , -0.23981026],
      [-0.27949443,  0.5497092 ,  0.22774768, -0.08467997]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.3807971 ,  0.        ,  0.4820138 ,  0.3807971 ,  0.        ],
       [-0.3807971 ,  0.3807971 ,  0.        ,  0.3807971 ,  0.4820138 ]],

      [[ 0.3807971 ,  0.23105858,  0.        ,  0.4820138 , -0.12245933],
       [ 0.4820138 , -0.4975274 , -0.49966466,  0.        ,  0.45257413]]]]

test_cc_attention_3d_with_past_and_present_qk_matmul_softmax

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    q_num_heads = 2
    kv_num_heads = 2
    qk_matmul_output_mode = 3
Inputs:
  Q: shape=(1, 2, 4), dtype=float32
    [[[ 1. ,  0. ,  0.5,  0.5],
      [ 0. ,  1. ,  1. , -1. ]]]
  K: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  , -1.  ,  1.  ],
      [ 0.5 ,  0.5 ,  1.  ,  1.  ],
      [ 0.  ,  1.  ,  0.25, -0.5 ]]]
  V: shape=(1, 3, 4), dtype=float32
    [[[ 1.  ,  0.  ,  2.  , -2.  ],
      [ 0.  ,  1.  ,  0.5 ,  0.25],
      [-1.  ,  1.  , -0.5 ,  0.  ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 4), dtype=float32
    [[[ 0.10764056,  0.45608166,  0.48057237, -0.20996696],
      [-0.31939295,  0.57814157,  0.05822894,  0.11373253]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[0.20710382, 0.14542593, 0.2949405 , 0.20710382, 0.14542593],
       [0.10673924, 0.21647945, 0.15200938, 0.21647945, 0.3082925 ]],

      [[0.21705809, 0.18188749, 0.15241571, 0.30911657, 0.13952214],
       [0.38144314, 0.06511761, 0.04572483, 0.18807767, 0.31963673]]]]

test_cc_attention_4d

Node:
  Attention(Q, K, V) -> (Y)
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.,  0.],
       [ 0.,  1.]],

      [[ 1.,  1.],
       [-1.,  1.]]]]
  V: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.,  2.],
       [ 3.,  4.]],

      [[-1.,  0.],
       [ 0.,  1.]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.6604769 ,  2.660477  ],
       [ 2.339523  ,  3.339523  ]],

      [[-0.66976154,  0.33023846],
       [-0.80442965,  0.19557032]]]]

test_cc_attention_4d_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.458196  ,  0.41020232],
       [ 0.0697852 ,  0.6150191 ]],

      [[ 0.9921614 , -0.7460807 ],
       [ 0.3994022 , -0.34422377]]]]

test_cc_attention_4d_attn_mask_3d

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 3), dtype=float32
    [[[ 0. , -1. ,  0.5],
      [ 0.2,  0. , -0.4]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.        ,  0.5626513 ],
       [ 0.03219185,  0.6617174 ]],

      [[ 0.48452467, -0.58578634],
       [ 0.3757273 , -0.26752764]]]]

test_cc_attention_4d_attn_mask_3d_causal

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [ 0.5 ,  0.5 ]],

      [[-1.  ,  1.  ],
       [ 1.  , -1.  ],
       [ 0.25,  0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 3, 3), dtype=float32
    [[[ 0. , -1. ,  0.5],
      [ 0.2,  0. , -0.4],
      [ 0.1, -0.3,  0. ]]]

Outputs:
  Y: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.        ,  0.        ],
       [ 0.4875026 ,  0.5124974 ],
       [ 0.03695408,  0.61167425]],

      [[ 2.        , -2.        ],
       [ 0.9650383 , -0.44755742],
       [ 0.7986912 , -0.689716  ]]]]

test_cc_attention_4d_attn_mask_4d

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0. , -0.5, -1. ],
       [ 0.2,  0. , -0.1]],

      [[-0.2,  0.3,  0. ],
       [ 0. , -0.1,  0.4]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.458196  ,  0.41020232],
       [-0.06765194,  0.6944395 ]],

      [[ 0.5724663 , -0.27137595],
       [ 0.02987868, -0.14798252]]]]

test_cc_attention_4d_attn_mask_4d_causal

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [ 0.5 ,  0.5 ]],

      [[-1.  ,  1.  ],
       [ 1.  , -1.  ],
       [ 0.25,  0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 0. , -0.5, -1. ],
       [ 0.2,  0. , -0.1],
       [ 0.5, -0.2,  0. ]],

      [[-0.2,  0.3,  0. ],
       [ 0. , -0.1,  0.4],
       [ 0.1,  0. , -0.3]]]]

Outputs:
  Y: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.        ,  0.        ],
       [ 0.4875026 ,  0.5124974 ],
       [ 0.18708876,  0.52451503]],

      [[ 2.        , -2.        ],
       [ 0.93357575, -0.40036362],
       [ 0.8560081 , -0.63305765]]]]

test_cc_attention_4d_attn_mask_bool

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=bool
    [[[[ True,  True, False],
       [ True, False,  True]],

      [[ True, False,  True],
       [ True,  True, False]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5621765 ,  0.4378235 ],
       [-0.24491866,  0.62245935]],

      [[ 0.7890498 , -1.0312399 ],
       [ 0.9034121 , -0.3551182 ]]]]

test_cc_attention_4d_attn_mask_bool_4d

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=bool
    [[[[ True,  True, False],
       [ True, False,  True]],

      [[ True, False,  True],
       [ True,  True, False]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5621765 ,  0.4378235 ],
       [-0.24491866,  0.62245935]],

      [[ 0.7890498 , -1.0312399 ],
       [ 0.9034121 , -0.3551182 ]]]]

test_cc_attention_4d_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [ 0.5 ,  0.5 ]],

      [[-1.  ,  1.  ],
       [ 1.  , -1.  ],
       [ 0.25,  0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.        ,  0.        ],
       [ 0.412521  ,  0.587479  ],
       [ 0.        ,  0.6666667 ]],

      [[ 2.        , -2.        ],
       [ 0.79335546, -0.19003323],
       [ 0.7425821 , -0.5205673 ]]]]

test_cc_attention_4d_causal_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.39085424,  0.15993309],
       [-0.01604789,  0.39012018]],

      [[ 0.7178145 , -0.5209762 ],
       [ 0.3204866 ,  0.16716442]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_diff_heads_mask4d_padded_kv

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 1, 2, 3), dtype=float32
    [[[[     0.,      0., -10000.],
       [     0.,      0., -10000.]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.49375033,  0.50624967],
       [ 0.49375033,  0.50624967]],

      [[ 0.4906261 ,  0.5093739 ],
       [ 0.53120935,  0.46879062]],

      [[ 1.066311  , -0.5994665 ],
       [ 1.25      , -0.875     ]],

      [[ 0.9034121 , -0.3551182 ],
       [ 1.066311  , -0.5994665 ]]]]

test_cc_attention_4d_diff_heads_sizes

Node:
  Attention(Q, K, V) -> (Y)
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.23092115,  0.5444725 ,  0.296508  ],
       [-0.23092115,  0.77539366,  0.64288974]],

      [[ 0.6482418 , -0.37858847,  0.3571369 ],
       [ 0.04638294, -0.08028026,  0.57520705]]]]

test_cc_attention_4d_diff_heads_sizes_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.458196  ,  0.41020232,  0.03320454],
       [ 0.0697852 ,  0.6150191 ,  0.3722638 ]],

      [[ 0.9921614 , -0.7460807 ,  0.4670852 ],
       [ 0.3994022 , -0.34422377,  0.5532167 ]]]]

test_cc_attention_4d_diff_heads_sizes_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [ 0.5 ,  0.5 ]],

      [[-1.  ,  1.  ],
       [ 1.  , -1.  ],
       [ 0.25,  0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]

Outputs:
  Y: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.        ,  0.        , -1.        ],
       [ 0.412521  ,  0.587479  ,  0.762437  ],
       [ 0.        ,  0.6666667 ,  0.5       ]],

      [[ 2.        , -2.        ,  1.        ],
       [ 0.79335546, -0.19003323, -0.0055371 ],
       [ 0.7425821 , -0.5205673 ,  0.43631145]]]]

test_cc_attention_4d_diff_heads_sizes_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.00333332,  0.6649986 ,  0.49749377],
       [-0.00333332,  0.6683319 ,  0.50249374]],

      [[ 0.6665943 , -0.58079195,  0.5803768 ],
       [ 0.65489024, -0.5724745 ,  0.5816217 ]]]]

test_cc_attention_4d_diff_heads_sizes_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 3), dtype=float32
    [[[[ 1.  ,  0.  , -1.  ],
       [ 0.  ,  1.  ,  2.  ],
       [-1.  ,  1.  ,  0.5 ]],

      [[ 2.  , -2.  ,  1.  ],
       [ 0.5 ,  0.25, -0.25],
       [-0.5 ,  0.  ,  1.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.15169363,  0.6033754 ,  0.44272968],
       [-0.15169363,  0.755069  ,  0.67027014]],

      [[ 0.67558366, -0.4552321 ,  0.4223395 ],
       [ 0.29168454, -0.3030643 ,  0.60680556]]]]

test_cc_attention_4d_diff_heads_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.12450287,  0.51090705],
       [-0.12746866,  0.51840895]],

      [[-0.13178737,  0.5070071 ],
       [ 0.00804363,  0.47238812]],

      [[ 0.48057237, -0.20996696],
       [ 0.7363033 , -0.5218315 ]],

      [[ 0.27224937,  0.02852735],
       [ 0.23465633, -0.06618155]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_diff_heads_with_past_and_present_mask3d

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 5), dtype=float32
    [[[ 0. , -0.5, -1. ,  0.2,  0. ],
      [ 0.5,  0. , -0.2, -0.1,  0. ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.19640994,  0.6580199 ],
       [-0.09580244,  0.52036375]],

      [[-0.19427064,  0.65238005],
       [ 0.00927786,  0.49179137]],

      [[ 0.30545416,  0.01247976],
       [ 0.560947  , -0.33633795]],

      [[ 0.20446573,  0.11752708],
       [ 0.21916693, -0.00962794]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_diff_heads_with_past_and_present_mask4d

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 1, 2, 5), dtype=float32
    [[[[ 0. , -0.5, -1. ,  0.2,  0. ],
       [ 0.5,  0. , -0.2, -0.1,  0. ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.19640994,  0.6580199 ],
       [-0.09580244,  0.52036375]],

      [[-0.19427064,  0.65238005],
       [ 0.00927786,  0.49179137]],

      [[ 0.30545416,  0.01247976],
       [ 0.560947  , -0.33633795]],

      [[ 0.20446573,  0.11752708],
       [ 0.21916693, -0.00962794]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_fp16

Node:
  Attention(Q, K, V) -> (Y)
Inputs:
  Q: shape=(2, 3, 4, 8), dtype=float16
    [[[[0.768   , 0.2769  , 0.381   , ..., 0.3386  , 0.9316  , 0.6865  ],
       [0.797   , 0.695   , 0.511   , ..., 0.941   , 0.3232  , 0.735   ],
       [0.0287  , 0.1525  , 0.3313  , ..., 0.9663  , 0.2722  , 0.742   ],
       [0.9644  , 0.906   , 0.924   , ..., 0.9604  , 0.619   , 0.382   ]],

      [[0.1196  , 0.4138  , 0.012245, ..., 0.9214  , 0.09467 , 0.907   ],
       [0.5264  , 0.84    , 0.4602  , ..., 0.802   , 0.821   , 0.6484  ],
       [0.232   , 0.549   , 0.772   , ..., 0.005775, 0.809   , 0.514   ],
       [0.7993  , 0.6104  , 0.593   , ..., 0.891   , 0.4631  , 0.4915  ]],

      [[0.8076  , 0.2961  , 0.208   , ..., 0.2695  , 0.07605 , 0.145   ],
       [0.35    , 0.581   , 0.0441  , ..., 0.9043  , 0.3318  , 0.1179  ],
       [0.537   , 0.6436  , 0.168   , ..., 0.01337 , 0.805   , 0.6313  ],
       [0.994   , 0.3652  , 0.787   , ..., 0.767   , 0.462   , 0.985   ]]],


     [[[0.3628  , 0.83    , 0.7495  , ..., 0.7886  , 0.1573  , 0.0558  ],
       [0.799   , 0.8257  , 0.04434 , ..., 0.655   , 0.1383  , 0.7993  ],
       [0.4753  , 0.3433  , 0.3005  , ..., 0.3953  , 0.542   , 0.75    ],
       [0.0801  , 0.0763  , 0.2883  , ..., 0.993   , 0.8965  , 0.5186  ]],

      [[0.3167  , 0.922   , 0.4185  , ..., 0.885   , 0.6196  , 0.296   ],
       [0.4043  , 0.48    , 0.2421  , ..., 0.4587  , 0.521   , 0.3496  ],
       [0.578   , 0.5537  , 0.951   , ..., 0.5576  , 0.3003  , 0.526   ],
       [0.588   , 0.65    , 0.878   , ..., 0.9766  , 0.04785 , 0.749   ]],

      [[0.2006  , 0.4778  , 0.9956  , ..., 0.965   , 0.7446  , 0.5215  ],
       [0.6973  , 0.4497  , 0.4187  , ..., 0.723   , 0.7617  , 0.4243  ],
       [0.376   , 0.1816  , 0.9004  , ..., 0.9067  , 0.3618  , 0.6157  ],
       [0.04916 , 0.992   , 0.3367  , ..., 0.12354 , 0.1976  , 0.3323  ]]]]
  K: shape=(2, 3, 6, 8), dtype=float16
    [[[[0.6216  , 0.861   , 0.5938  , ..., 0.0877  , 0.545   , 0.1114  ],
       [0.04868 , 0.9824  , 0.7256  , ..., 0.923   , 0.00809 , 0.6504  ],
       [0.9204  , 0.1576  , 0.1031  , ..., 0.1365  , 0.8228  , 0.3477  ],
       [0.4634  , 0.08734 , 0.598   , ..., 0.8447  , 0.2439  , 0.26    ],
       [0.171   , 0.927   , 0.973   , ..., 0.8657  , 0.666   , 0.2727  ],
       [0.3013  , 0.7837  , 0.381   , ..., 0.222   , 0.1247  , 0.599   ]],

      [[0.378   , 0.2192  , 0.8647  , ..., 0.6387  , 0.7964  , 0.281   ],
       [0.691   , 0.7476  , 0.8584  , ..., 0.04382 , 0.4983  , 0.688   ],
       [0.795   , 0.3357  , 0.6875  , ..., 0.0688  , 0.1875  , 0.0173  ],
       [0.010956, 0.9043  , 0.0662  , ..., 0.548   , 0.4624  , 0.7954  ],
       [0.926   , 0.827   , 0.601   , ..., 0.01913 , 0.06143 , 0.376   ],
       [0.8906  , 0.4302  , 0.289   , ..., 0.2229  , 0.3643  , 0.4412  ]],

      [[0.5244  , 0.494   , 0.01675 , ..., 0.603   , 0.891   , 0.5015  ],
       [0.2092  , 0.751   , 0.06223 , ..., 0.8066  , 0.8022  , 0.395   ],
       [0.5947  , 0.3872  , 0.5938  , ..., 0.184   , 0.399   , 0.788   ],
       [0.09515 , 0.542   , 0.1709  , ..., 0.2878  , 0.623   , 0.6763  ],
       [0.3906  , 0.8076  , 0.2429  , ..., 0.375   , 0.5576  , 0.8813  ],
       [0.9272  , 0.7295  , 0.6504  , ..., 0.905   , 0.9253  , 0.554   ]]],


     [[[0.4163  , 0.3062  , 0.1239  , ..., 0.03952 , 0.59    , 0.675   ],
       [0.3347  , 0.489   , 0.7856  , ..., 0.1453  , 0.05606 , 0.5493  ],
       [0.4255  , 0.683   , 0.647   , ..., 0.2908  , 0.9688  , 0.313   ],
       [0.1964  , 0.2465  , 0.1116  , ..., 0.748   , 0.6133  , 0.43    ],
       [0.4219  , 0.9897  , 0.4724  , ..., 0.492   , 0.416   , 0.191   ],
       [0.641   , 0.6777  , 0.4138  , ..., 0.701   , 0.4443  , 0.2162  ]],

      [[0.6597  , 0.527   , 0.5103  , ..., 0.2566  , 0.9053  , 0.953   ],
       [0.04782 , 0.708   , 0.7275  , ..., 0.2428  , 0.647   , 0.1777  ],
       [0.6426  , 0.8447  , 0.921   , ..., 0.4475  , 0.6587  , 0.494   ],
       [0.04538 , 0.5137  , 0.3374  , ..., 0.8613  , 0.5693  , 0.834   ],
       [0.1244  , 0.7446  , 0.11426 , ..., 0.1786  , 0.1487  , 0.9575  ],
       [0.5127  , 0.5205  , 0.7793  , ..., 0.2957  , 0.808   , 0.954   ]],

      [[0.1094  , 0.277   , 0.7505  , ..., 0.8125  , 0.09845 , 0.7393  ],
       [0.5786  , 0.4038  , 0.8374  , ..., 0.533   , 0.2117  , 0.558   ],
       [0.851   , 0.7427  , 0.7397  , ..., 0.9634  , 0.4805  , 0.4832  ],
       [0.6226  , 0.0897  , 0.548   , ..., 0.8125  , 0.8774  , 0.915   ],
       [0.854   , 0.693   , 0.2424  , ..., 0.4934  , 0.517   , 0.0836  ],
       [0.2725  , 0.754   , 0.1952  , ..., 0.6216  , 0.1534  , 0.5454  ]]]]
  V: shape=(2, 3, 6, 8), dtype=float16
    [[[[0.4753  , 0.4448  , 0.806   , ..., 0.837   , 0.1581  , 0.5366  ],
       [0.3005  , 0.2703  , 0.9395  , ..., 0.905   , 0.693   , 0.566   ],
       [0.812   , 0.1626  , 0.875   , ..., 0.3066  , 0.3735  , 0.953   ],
       [0.963   , 0.269   , 0.273   , ..., 0.7285  , 0.8687  , 0.1381  ],
       [0.2224  , 0.4397  , 0.9346  , ..., 0.8096  , 0.2372  , 0.638   ],
       [0.076   , 0.728   , 0.302   , ..., 0.642   , 0.4287  , 0.5503  ]],

      [[0.524   , 0.8896  , 0.9575  , ..., 0.2717  , 0.7837  , 0.04794 ],
       [0.582   , 0.885   , 0.1246  , ..., 0.1964  , 0.533   , 0.8843  ],
       [0.782   , 0.3752  , 0.1667  , ..., 0.868   , 0.2988  , 0.8896  ],
       [0.672   , 0.2277  , 0.0884  , ..., 0.1917  , 0.593   , 0.4727  ],
       [0.314   , 0.01035 , 0.03458 , ..., 0.0249  , 0.3174  , 0.1203  ],
       [0.787   , 0.4954  , 0.791   , ..., 0.6787  , 0.2668  , 0.8975  ]],

      [[0.686   , 0.1578  , 0.284   , ..., 0.418   , 0.6245  , 0.947   ],
       [0.6196  , 0.676   , 0.0801  , ..., 0.959   , 0.4658  , 0.99    ],
       [0.714   , 0.4312  , 0.887   , ..., 0.9727  , 0.2556  , 0.8267  ],
       [0.11017 , 0.007835, 0.05338 , ..., 0.5825  , 0.3499  , 0.8335  ],
       [0.4648  , 0.6934  , 0.0673  , ..., 0.865   , 0.4956  , 0.4666  ],
       [0.4504  , 0.979   , 0.05844 , ..., 0.3506  , 0.3293  , 0.7593  ]]],


     [[[0.255   , 0.05832 , 0.2969  , ..., 0.5215  , 0.8794  , 0.823   ],
       [0.08167 , 0.3281  , 0.693   , ..., 0.3137  , 0.0643  , 0.349   ],
       [0.6504  , 0.8887  , 0.2983  , ..., 0.6167  , 0.1927  , 0.10443 ],
       [0.6953  , 0.04318 , 0.8047  , ..., 0.773   , 0.4648  , 0.4355  ],
       [0.4678  , 0.798   , 0.04443 , ..., 0.07715 , 0.4705  , 0.7666  ],
       [0.2332  , 0.3625  , 0.491   , ..., 0.2786  , 0.691   , 0.1003  ]],

      [[0.7734  , 0.6494  , 0.758   , ..., 0.0784  , 0.9976  , 0.517   ],
       [0.886   , 0.7744  , 0.0998  , ..., 0.632   , 0.1531  , 0.6753  ],
       [0.884   , 0.5566  , 0.9116  , ..., 0.04623 , 0.3098  , 0.3125  ],
       [0.09576 , 0.01733 , 0.2286  , ..., 0.8833  , 0.511   , 0.743   ],
       [0.3657  , 0.8823  , 0.2277  , ..., 0.6562  , 0.1917  , 0.3606  ],
       [0.5537  , 0.07935 , 0.725   , ..., 0.3335  , 0.676   , 0.1355  ]],

      [[0.0349  , 0.24    , 0.679   , ..., 0.835   , 0.6787  , 0.1176  ],
       [0.2008  , 0.1987  , 0.6875  , ..., 0.5337  , 0.806   , 0.934   ],
       [0.958   , 0.4932  , 0.4885  , ..., 0.7573  , 0.05563 , 0.291   ],
       [0.2291  , 0.864   , 0.4624  , ..., 0.285   , 0.3137  , 0.4717  ],
       [0.4517  , 0.755   , 0.1288  , ..., 0.8496  , 0.858   , 0.838   ],
       [0.07965 , 0.813   , 0.6484  , ..., 0.637   , 0.8574  , 0.3298  ]]]]

Outputs:
  Y: shape=(2, 3, 4, 8), dtype=float16
    [[[[0.4634, 0.3958, 0.698 , ..., 0.7026, 0.4287, 0.576 ],
       [0.449 , 0.3945, 0.6997, ..., 0.725 , 0.4514, 0.5586],
       [0.458 , 0.394 , 0.684 , ..., 0.7275, 0.466 , 0.5444],
       [0.4526, 0.3914, 0.7095, ..., 0.7305, 0.4438, 0.5566]],

      [[0.6167, 0.5005, 0.366 , ..., 0.3538, 0.4849, 0.5576],
       [0.6084, 0.5273, 0.372 , ..., 0.3508, 0.4912, 0.5522],
       [0.613 , 0.5356, 0.3718, ..., 0.3594, 0.4858, 0.5713],
       [0.6074, 0.502 , 0.3713, ..., 0.3584, 0.4749, 0.552 ]],

      [[0.509 , 0.5264, 0.2328, ..., 0.679 , 0.4148, 0.7964],
       [0.512 , 0.531 , 0.217 , ..., 0.6816, 0.4224, 0.8027],
       [0.5093, 0.5264, 0.228 , ..., 0.6885, 0.4185, 0.792 ],
       [0.5093, 0.5566, 0.2291, ..., 0.671 , 0.4097, 0.788 ]]],


     [[[0.4038, 0.4631, 0.42  , ..., 0.4077, 0.4414, 0.405 ],
       [0.4033, 0.448 , 0.425 , ..., 0.414 , 0.4512, 0.4102],
       [0.4038, 0.4336, 0.4287, ..., 0.429 , 0.4556, 0.418 ],
       [0.4204, 0.4429, 0.4329, ..., 0.4365, 0.4502, 0.4038]],

      [[0.5933, 0.4773, 0.4858, ..., 0.4468, 0.4697, 0.4766],
       [0.5947, 0.4795, 0.498 , ..., 0.4353, 0.4797, 0.4631],
       [0.6094, 0.4758, 0.5244, ..., 0.4106, 0.4888, 0.4495],
       [0.5913, 0.4678, 0.509 , ..., 0.4304, 0.4807, 0.4573]],

      [[0.358 , 0.552 , 0.5205, ..., 0.6343, 0.546 , 0.4905],
       [0.3645, 0.5645, 0.5103, ..., 0.6333, 0.5483, 0.5015],
       [0.353 , 0.5474, 0.5225, ..., 0.6353, 0.552 , 0.4941],
       [0.348 , 0.556 , 0.515 , ..., 0.651 , 0.5835, 0.5044]]]]

test_cc_attention_4d_gqa

Node:
  Attention(Q, K, V) -> (Y)
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.02356532,  0.6783799 ],
       [-0.02356531,  0.6783799 ]],

      [[-0.03533878,  0.6841799 ],
       [ 0.11724145,  0.6063233 ]],

      [[ 0.6482418 , -0.37858847],
       [ 0.9917567 , -0.74587834]],

      [[ 0.37784207, -0.12898168],
       [ 0.29831943, -0.26321504]]]]

test_cc_attention_4d_gqa_attn_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.30531266,  0.50214726],
       [ 0.2229336 ,  0.5335671 ]],

      [[ 0.29782698,  0.50647473],
       [ 0.31907293,  0.4797093 ]],

      [[ 0.9921614 , -0.7460807 ],
       [ 1.1873223 , -1.0039468 ]],

      [[ 0.74361753, -0.47596362],
       [ 0.6715176 , -0.5857588 ]]]]

test_cc_attention_4d_gqa_causal

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    is_causal = 1
Inputs:
  Q: shape=(1, 4, 3, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ],
       [-0.1 ,  0.05]],

      [[ 0.2 , -0.3 ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ],
       [ 0.25,  0.1 ]],

      [[-0.5 ,  0.5 ],
       [-0.25,  0.75],
       [ 0.1 , -0.1 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 4, 3, 2), dtype=float32
    [[[[ 1.        ,  0.        ],
       [ 0.4911621 ,  0.50883794],
       [-0.03533878,  0.6841799 ]],

      [[ 1.        ,  0.        ],
       [ 0.5       ,  0.5       ],
       [-0.23092115,  0.77539366]],

      [[ 2.        , -2.        ],
       [ 0.9953577 , -0.4930365 ],
       [ 0.60666823, -0.46362755]],

      [[ 2.        , -2.        ],
       [ 1.3812186 , -1.0718278 ],
       [ 0.5847607 , -0.50853795]]]]

test_cc_attention_4d_gqa_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-3.3333333e-04,  6.6683334e-01],
       [-3.3333330e-04,  6.6683334e-01]],

      [[-4.9999997e-04,  6.6691661e-01],
       [ 1.6666650e-03,  6.6583300e-01]],

      [[ 6.6659433e-01, -5.8079195e-01],
       [ 6.7248535e-01, -5.8624268e-01]],

      [[ 6.6071773e-01, -5.7536310e-01],
       [ 6.6077113e-01, -5.7789290e-01]]]]

test_cc_attention_4d_gqa_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.03040495,  0.6819007 ],
       [-0.02114749,  0.67739695]],

      [[-0.04947086,  0.69103485],
       [ 0.15203254,  0.5866078 ]],

      [[ 0.67558366, -0.4552321 ],
       [ 0.94510204, -0.72255105]],

      [[ 0.40459138, -0.23742941],
       [ 0.35040644, -0.32931072]]]]

test_cc_attention_4d_gqa_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
Inputs:
  Q: shape=(1, 4, 2, 2), dtype=float32
    [[[[ 0.1 ,  0.2 ],
       [ 0.3 ,  0.4 ]],

      [[-0.1 ,  0.05],
       [ 0.2 , -0.3 ]],

      [[ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [ 0.5 , -0.5 ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 4, 2, 2), dtype=float32
    [[[[-0.12450287,  0.51090705],
       [-0.12746866,  0.51840895]],

      [[-0.13178737,  0.5070071 ],
       [ 0.00804363,  0.47238812]],

      [[ 0.48057237, -0.20996696],
       [ 0.7363033 , -0.5218315 ]],

      [[ 0.27224937,  0.02852735],
       [ 0.23465633, -0.06618155]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_gqa_with_past_and_present_fp16

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value)
Inputs:
  Q: shape=(2, 9, 4, 8), dtype=float16
    [[[[0.0945  , 0.7554  , 0.3372  , ..., 0.2251  , 0.09625 , 0.5435  ],
       [0.3733  , 0.892   , 0.391   , ..., 0.666   , 0.8926  , 0.8906  ],
       [0.4827  , 0.1914  , 0.3293  , ..., 0.4897  , 0.82    , 0.8774  ],
       [0.995   , 0.8486  , 0.4473  , ..., 0.9634  , 0.0887  , 0.977   ]],

      [[0.8115  , 0.326   , 0.4934  , ..., 0.8677  , 0.611   , 0.8037  ],
       [0.661   , 0.8496  , 0.1693  , ..., 0.001018, 0.5024  , 0.6113  ],
       [0.6006  , 0.9126  , 0.63    , ..., 0.676   , 0.5825  , 0.7954  ],
       [0.5156  , 0.7734  , 0.983   , ..., 0.2205  , 0.8726  , 0.3906  ]],

      [[0.6196  , 0.9556  , 0.7896  , ..., 0.4795  , 0.0976  , 0.574   ],
       [0.9536  , 0.7485  , 0.564   , ..., 0.313   , 0.1863  , 0.1627  ],
       [0.8877  , 0.708   , 0.2732  , ..., 0.0961  , 0.77    , 0.11365 ],
       [0.6196  , 0.1549  , 0.838   , ..., 0.2201  , 0.6826  , 0.02515 ]],

      ...,

      [[0.8877  , 0.687   , 0.4976  , ..., 0.2415  , 0.888   , 0.3813  ],
       [0.3013  , 0.3867  , 0.4314  , ..., 0.954   , 0.781   , 0.2961  ],
       [0.854   , 0.2605  , 0.4016  , ..., 0.749   , 0.895   , 0.01886 ],
       [0.993   , 0.8794  , 0.579   , ..., 0.2678  , 0.315   , 0.3975  ]],

      [[0.4949  , 0.07935 , 0.587   , ..., 0.667   , 0.828   , 0.4207  ],
       [0.9634  , 0.4626  , 0.001799, ..., 0.1188  , 0.4258  , 0.718   ],
       [0.3308  , 0.8975  , 0.854   , ..., 0.3113  , 0.8013  , 0.05927 ],
       [0.357   , 0.01604 , 0.1256  , ..., 0.948   , 0.852   , 0.8555  ]],

      [[0.1307  , 0.7183  , 0.2534  , ..., 0.2489  , 0.1777  , 0.657   ],
       [0.5684  , 0.668   , 0.6265  , ..., 0.4478  , 0.5815  , 0.657   ],
       [0.9146  , 0.796   , 0.0865  , ..., 0.2957  , 0.5703  , 0.1877  ],
       [0.2417  , 0.7974  , 0.4294  , ..., 0.5586  , 0.853   , 0.2217  ]]],


     [[[0.951   , 0.6343  , 0.904   , ..., 0.5186  , 0.843   , 0.873   ],
       [0.272   , 0.0337  , 0.711   , ..., 0.4724  , 0.1549  , 0.896   ],
       [0.7607  , 0.988   , 0.506   , ..., 0.2335  , 0.1082  , 0.1857  ],
       [0.803   , 0.2551  , 0.8965  , ..., 0.6304  , 0.2249  , 0.2778  ]],

      [[0.1129  , 0.86    , 0.944   , ..., 0.007214, 0.73    , 0.8486  ],
       [0.2307  , 0.591   , 0.662   , ..., 0.724   , 0.05148 , 0.2148  ],
       [0.02669 , 0.5044  , 0.5186  , ..., 0.656   , 0.321   , 0.8345  ],
       [0.1984  , 0.42    , 0.933   , ..., 0.6943  , 0.873   , 0.8047  ]],

      [[0.772   , 0.425   , 0.7803  , ..., 0.2461  , 0.745   , 0.3655  ],
       [0.601   , 0.371   , 0.8857  , ..., 0.7334  , 0.1779  , 0.3953  ],
       [0.8677  , 0.3762  , 0.5303  , ..., 0.0942  , 0.1152  , 0.9146  ],
       [0.583   , 0.3235  , 0.946   , ..., 0.2822  , 0.704   , 0.584   ]],

      ...,

      [[0.506   , 0.712   , 0.2119  , ..., 0.3293  , 0.2125  , 0.4702  ],
       [0.3206  , 0.346   , 0.6333  , ..., 0.977   , 0.1055  , 0.358   ],
       [0.11206 , 0.04666 , 0.695   , ..., 0.631   , 0.2296  , 0.1803  ],
       [0.1407  , 0.9473  , 0.8794  , ..., 0.4949  , 0.982   , 0.848   ]],

      [[0.995   , 0.4465  , 0.12274 , ..., 0.288   , 0.4624  , 0.412   ],
       [0.0913  , 0.7095  , 0.3145  , ..., 0.7456  , 0.974   , 0.565   ],
       [0.1742  , 0.166   , 0.797   , ..., 0.1173  , 0.524   , 0.1388  ],
       [0.8164  , 0.01196 , 0.8647  , ..., 0.67    , 0.321   , 0.607   ]],

      [[0.919   , 0.7085  , 0.267   , ..., 0.1848  , 0.2778  , 0.583   ],
       [0.2109  , 0.4834  , 0.7056  , ..., 0.4595  , 0.5103  , 0.321   ],
       [0.7495  , 0.8286  , 0.3352  , ..., 0.807   , 0.8374  , 0.2168  ],
       [0.4192  , 0.002748, 0.2634  , ..., 0.556   , 0.2805  , 0.3054  ]]]]
  K: shape=(2, 3, 6, 8), dtype=float16
    [[[[8.8623e-01, 7.0361e-01, 8.4521e-01, ..., 2.0764e-01, 7.2363e-01,
        8.4766e-01],
       [6.6455e-01, 3.1860e-01, 8.8867e-01, ..., 3.4546e-01, 9.4727e-01,
        5.0342e-01],
       [5.1367e-01, 6.5186e-01, 7.2949e-01, ..., 4.2432e-01, 2.4097e-01,
        9.3213e-01],
       [6.8115e-01, 6.8457e-01, 8.6914e-01, ..., 5.5811e-01, 8.3301e-01,
        8.6084e-01],
       [3.6938e-01, 7.8760e-01, 8.8867e-01, ..., 5.0098e-01, 2.8125e-01,
        2.8174e-01],
       [2.3267e-01, 2.2388e-01, 9.4727e-01, ..., 1.4807e-01, 9.7021e-01,
        9.5312e-01]],

      [[8.7939e-01, 6.4551e-01, 2.8125e-01, ..., 3.5065e-02, 6.1377e-01,
        8.9990e-01],
       [3.6987e-01, 5.9668e-01, 7.0557e-01, ..., 8.3838e-01, 7.5293e-01,
        9.1016e-01],
       [7.1875e-01, 1.0910e-02, 1.1322e-01, ..., 8.7451e-01, 2.5781e-01,
        4.0222e-02],
       [3.9355e-01, 7.1338e-01, 9.7559e-01, ..., 6.0107e-01, 8.2617e-01,
        1.1041e-01],
       [8.1445e-01, 9.1895e-01, 8.4229e-01, ..., 1.1676e-01, 4.8340e-01,
        2.0715e-01],
       [8.5449e-01, 7.3340e-01, 8.3984e-01, ..., 6.5967e-01, 8.4045e-02,
        1.8262e-01]],

      [[3.4131e-01, 6.5430e-01, 1.7468e-01, ..., 1.0962e-01, 8.1006e-01,
        1.5454e-01],
       [5.5371e-01, 6.7993e-02, 6.3037e-01, ..., 4.8633e-01, 6.7139e-01,
        5.9700e-03],
       [7.2510e-01, 7.5244e-01, 6.8237e-02, ..., 4.5068e-01, 6.8932e-03,
        8.8623e-01],
       [5.4199e-01, 3.7646e-01, 9.2871e-01, ..., 8.0371e-01, 9.8291e-01,
        7.0898e-01],
       [6.4258e-01, 9.9902e-01, 7.2949e-01, ..., 9.8779e-01, 9.3750e-02,
        6.5576e-01],
       [1.1945e-01, 5.6934e-01, 5.6738e-01, ..., 5.8447e-01, 6.6260e-01,
        6.7139e-01]]],


     [[[1.8356e-02, 4.2847e-01, 6.1572e-01, ..., 8.1787e-01, 3.4009e-01,
        9.6729e-01],
       [8.3740e-01, 8.0713e-01, 6.2793e-01, ..., 5.1575e-02, 6.0547e-01,
        5.2100e-01],
       [2.8870e-02, 3.2642e-01, 4.3408e-01, ..., 2.8955e-01, 2.6562e-01,
        5.7520e-01],
       [4.9683e-01, 4.9902e-01, 4.4312e-01, ..., 6.7725e-01, 9.5557e-01,
        1.3770e-01],
       [9.9731e-02, 2.2778e-01, 1.4148e-01, ..., 9.8050e-05, 6.3965e-01,
        4.8267e-01],
       [1.4832e-01, 3.0591e-01, 5.9424e-01, ..., 1.8433e-01, 1.0815e-01,
        6.4990e-01]],

      [[9.0674e-01, 9.1748e-01, 9.4482e-01, ..., 7.9688e-01, 4.8145e-01,
        9.4385e-01],
       [9.2407e-02, 1.3782e-01, 9.1357e-01, ..., 5.4541e-01, 7.0703e-01,
        4.3628e-01],
       [3.7573e-01, 4.3140e-01, 3.0103e-01, ..., 7.7832e-01, 6.1127e-02,
        4.6313e-01],
       [8.7988e-01, 1.5479e-01, 4.8364e-01, ..., 9.8340e-01, 6.4746e-01,
        1.2622e-01],
       [6.8213e-01, 5.4199e-02, 9.1748e-01, ..., 2.9102e-01, 8.9795e-01,
        7.9346e-01],
       [3.1226e-01, 7.6709e-01, 6.3574e-01, ..., 9.7070e-01, 7.3120e-02,
        5.9863e-01]],

      [[2.5244e-01, 3.2104e-01, 7.5098e-01, ..., 9.3311e-01, 7.6123e-01,
        9.3994e-01],
       [4.3945e-01, 6.3477e-01, 9.5215e-01, ..., 2.2998e-01, 3.7915e-01,
        9.8193e-01],
       [7.6172e-01, 1.7120e-02, 8.3466e-03, ..., 5.5273e-01, 6.7090e-01,
        1.5527e-01],
       [5.6250e-01, 6.6797e-01, 2.6514e-01, ..., 7.3389e-01, 2.0972e-01,
        6.5527e-01],
       [1.3635e-01, 1.7761e-01, 1.4709e-01, ..., 2.4695e-01, 8.9648e-01,
        9.4385e-01],
       [2.3218e-01, 2.6993e-02, 6.5674e-01, ..., 7.0557e-01, 4.5972e-01,
        2.4805e-01]]]]
  V: shape=(2, 3, 6, 8), dtype=float16
    [[[[0.2004  , 0.4282  , 0.9204  , ..., 0.06033 , 0.3813  , 0.8096  ],
       [0.1216  , 0.3513  , 0.8833  , ..., 0.4768  , 0.5493  , 0.9263  ],
       [0.4238  , 0.597   , 0.812   , ..., 0.499   , 0.743   , 0.3118  ],
       [0.3052  , 0.486   , 0.628   , ..., 0.5186  , 0.7974  , 0.564   ],
       [0.291   , 0.605   , 0.704   , ..., 0.442   , 0.25    , 0.422   ],
       [0.736   , 0.3054  , 0.3667  , ..., 0.1913  , 0.8413  , 0.265   ]],

      [[0.3218  , 0.2054  , 0.3962  , ..., 0.205   , 0.01563 , 0.613   ],
       [0.5586  , 0.6885  , 0.5254  , ..., 0.9365  , 0.4194  , 0.6265  ],
       [0.2852  , 0.4028  , 0.94    , ..., 0.3547  , 0.403   , 0.607   ],
       [0.1678  , 0.7637  , 0.7793  , ..., 0.445   , 0.52    , 0.4966  ],
       [0.2012  , 0.451   , 0.6357  , ..., 0.708   , 0.02582 , 0.3772  ],
       [0.208   , 0.91    , 0.05463 , ..., 0.659   , 0.3806  , 0.9727  ]],

      [[0.339   , 0.8535  , 0.03452 , ..., 0.3306  , 0.747   , 0.1459  ],
       [0.5737  , 0.7573  , 0.527   , ..., 0.7695  , 0.4905  , 0.402   ],
       [0.2186  , 0.3645  , 0.937   , ..., 0.539   , 0.6797  , 0.3857  ],
       [0.4094  , 0.1833  , 0.6226  , ..., 0.218   , 0.5874  , 0.383   ],
       [0.6094  , 0.988   , 0.394   , ..., 0.9004  , 0.1876  , 0.82    ],
       [0.1097  , 0.3174  , 0.01627 , ..., 0.696   , 0.6597  , 0.7627  ]]],


     [[[0.5303  , 0.9775  , 0.7065  , ..., 0.2302  , 0.884   , 0.4202  ],
       [0.319   , 0.5386  , 0.635   , ..., 0.1444  , 0.9453  , 0.6396  ],
       [0.2515  , 0.8374  , 0.4248  , ..., 0.595   , 0.6313  , 0.4097  ],
       [0.932   , 0.4756  , 0.6533  , ..., 0.7554  , 0.933   , 0.3604  ],
       [0.2925  , 0.3213  , 0.3499  , ..., 0.312   , 0.544   , 0.2612  ],
       [0.09314 , 0.007397, 0.4978  , ..., 0.969   , 0.3613  , 0.5234  ]],

      [[0.4224  , 0.433   , 0.5327  , ..., 0.947   , 0.985   , 0.6973  ],
       [0.1969  , 0.263   , 0.844   , ..., 0.4788  , 0.218   , 0.9756  ],
       [0.661   , 0.427   , 0.274   , ..., 0.316   , 0.567   , 0.6904  ],
       [0.888   , 0.621   , 0.618   , ..., 0.2245  , 0.7417  , 0.8154  ],
       [0.278   , 0.8066  , 0.12445 , ..., 0.486   , 0.1539  , 0.4636  ],
       [0.0608  , 0.7104  , 0.995   , ..., 0.398   , 0.42    , 0.3901  ]],

      [[0.793   , 0.326   , 0.885   , ..., 0.7734  , 0.859   , 0.4893  ],
       [0.86    , 0.4106  , 0.4014  , ..., 0.4407  , 0.992   , 0.297   ],
       [0.975   , 0.989   , 0.606   , ..., 0.744   , 0.04483 , 0.49    ],
       [0.6797  , 0.85    , 0.5127  , ..., 0.04388 , 0.4456  , 0.3848  ],
       [0.8438  , 0.04922 , 0.589   , ..., 0.2152  , 0.3254  , 0.4387  ],
       [0.165   , 0.4075  , 0.2546  , ..., 0.1492  , 0.01906 , 0.75    ]]]]
  attn_mask: shape=(4, 18), dtype=float16
    [[0.1581  , 0.3494  , 0.621   , ..., 0.509   , 0.999   , 0.0327  ],
     [0.4724  , 0.548   , 0.07904 , ..., 0.764   , 0.4182  , 0.7754  ],
     [0.39    , 0.1998  , 0.6245  , ..., 0.004356, 0.9424  , 0.09125 ],
     [0.673   , 0.6772  , 0.1316  , ..., 0.4263  , 0.4575  , 0.7085  ]]
  past_key: shape=(2, 3, 12, 8), dtype=float16
    [[[[0.5303  , 0.1553  , 0.91    , ..., 0.3093  , 0.8994  , 0.2177  ],
       [0.3384  , 0.662   , 0.66    , ..., 0.573   , 0.0795  , 0.928   ],
       [0.534   , 0.2969  , 0.188   , ..., 0.6865  , 0.2976  , 0.03244 ],
       ...,
       [0.9775  , 0.08466 , 0.649   , ..., 0.0716  , 0.2751  , 0.7896  ],
       [0.8716  , 0.4856  , 0.07556 , ..., 0.383   , 0.003033, 0.71    ],
       [0.912   , 0.622   , 0.933   , ..., 0.0661  , 0.3047  , 0.7373  ]],

      [[0.775   , 0.03592 , 0.011   , ..., 0.01388 , 0.54    , 0.3003  ],
       [0.777   , 0.5967  , 0.3645  , ..., 0.6978  , 0.5845  , 0.53    ],
       [0.3755  , 0.001608, 0.8145  , ..., 0.6665  , 0.3289  , 0.7603  ],
       ...,
       [0.2415  , 0.002747, 0.394   , ..., 0.9844  , 0.2812  , 0.1603  ],
       [0.9604  , 0.6587  , 0.03964 , ..., 0.3857  , 0.5825  , 0.8164  ],
       [0.2223  , 0.2042  , 0.1219  , ..., 0.03165 , 0.807   , 0.1245  ]],

      [[0.3904  , 0.1189  , 0.618   , ..., 0.753   , 0.503   , 0.9507  ],
       [0.4692  , 0.2097  , 0.7705  , ..., 0.958   , 0.732   , 0.8633  ],
       [0.604   , 0.1118  , 0.588   , ..., 0.1339  , 0.573   , 0.6343  ],
       ...,
       [0.8223  , 0.4124  , 0.403   , ..., 0.846   , 0.493   , 0.575   ],
       [0.499   , 0.04944 , 0.713   , ..., 0.6885  , 0.938   , 0.381   ],
       [0.3647  , 0.981   , 0.7153  , ..., 0.899   , 0.906   , 0.634   ]]],


     [[[0.4712  , 0.623   , 0.5747  , ..., 0.245   , 0.1318  , 0.637   ],
       [0.51    , 0.7207  , 0.7217  , ..., 0.07306 , 0.8657  , 0.3967  ],
       [0.313   , 0.846   , 0.3518  , ..., 0.8076  , 0.3745  , 0.1225  ],
       ...,
       [0.8135  , 0.532   , 0.52    , ..., 0.3748  , 0.2534  , 0.503   ],
       [0.5806  , 0.877   , 0.939   , ..., 0.010124, 0.4126  , 0.8726  ],
       [0.4336  , 0.1708  , 0.557   , ..., 0.3865  , 0.9443  , 0.7344  ]],

      [[0.1107  , 0.767   , 0.725   , ..., 0.209   , 0.771   , 0.8286  ],
       [0.9917  , 0.348   , 0.0619  , ..., 0.7617  , 0.3303  , 0.5996  ],
       [0.5967  , 0.4238  , 0.95    , ..., 0.4062  , 0.0765  , 0.694   ],
       ...,
       [0.3086  , 0.581   , 0.588   , ..., 0.2893  , 0.9067  , 0.4124  ],
       [0.299   , 0.3604  , 0.562   , ..., 0.0693  , 0.3503  , 0.7603  ],
       [0.522   , 0.9927  , 0.4907  , ..., 0.213   , 0.266   , 0.895   ]],

      [[0.4028  , 0.7695  , 0.913   , ..., 0.3638  , 0.7637  , 0.9946  ],
       [0.00923 , 0.3105  , 0.635   , ..., 0.7427  , 0.469   , 0.9404  ],
       [0.5483  , 0.063   , 0.2261  , ..., 0.649   , 0.0228  , 0.8174  ],
       ...,
       [0.4016  , 0.7773  , 0.4514  , ..., 0.3088  , 0.797   , 0.7725  ],
       [0.7485  , 0.718   , 0.4255  , ..., 0.5845  , 0.09503 , 0.512   ],
       [0.7236  , 0.6655  , 0.36    , ..., 0.09753 , 0.215   , 0.5835  ]]]]
  past_value: shape=(2, 3, 12, 8), dtype=float16
    [[[[0.608  , 0.9365 , 0.646  , ..., 0.83   , 0.3618 , 0.2002 ],
       [0.1482 , 0.741  , 0.8584 , ..., 0.368  , 0.3079 , 0.9976 ],
       [0.6387 , 0.824  , 0.777  , ..., 0.1534 , 0.9316 , 0.599  ],
       ...,
       [0.608  , 0.4153 , 0.738  , ..., 0.949  , 0.0758 , 0.6973 ],
       [0.05075, 0.3218 , 0.04764, ..., 0.5913 , 0.5225 , 0.9336 ],
       [0.329  , 0.7646 , 0.171  , ..., 0.8096 , 0.669  , 0.00903]],

      [[0.1675 , 0.921  , 0.1498 , ..., 0.0378 , 0.908  , 0.7837 ],
       [0.4946 , 0.6074 , 0.3533 , ..., 0.601  , 0.2115 , 0.6333 ],
       [0.9414 , 0.781  , 0.9795 , ..., 0.2142 , 0.6284 , 0.9487 ],
       ...,
       [0.11365, 0.607  , 0.717  , ..., 0.0315 , 0.268  , 0.3342 ],
       [0.6226 , 0.02585, 0.2186 , ..., 0.6733 , 0.742  , 0.3315 ],
       [0.0741 , 0.378  , 0.6885 , ..., 0.8813 , 0.7886 , 0.5176 ]],

      [[0.631  , 0.527  , 0.1058 , ..., 0.7773 , 0.739  , 0.94   ],
       [0.658  , 0.9775 , 0.7773 , ..., 0.373  , 0.5015 , 0.1617 ],
       [0.2258 , 0.3743 , 0.838  , ..., 0.072  , 0.06244, 0.7617 ],
       ...,
       [0.7495 , 0.3655 , 0.4368 , ..., 0.7153 , 0.6587 , 0.4268 ],
       [0.4885 , 0.4294 , 0.8574 , ..., 0.6147 , 0.5117 , 0.755  ],
       [0.5273 , 0.07324, 0.651  , ..., 0.3203 , 0.06027, 0.872  ]]],


     [[[0.4666 , 0.848  , 0.733  , ..., 0.6416 , 0.5728 , 0.01241],
       [0.1087 , 0.946  , 0.3486 , ..., 0.2783 , 0.8955 , 0.927  ],
       [0.961  , 0.6987 , 0.5703 , ..., 0.3208 , 0.9526 , 0.381  ],
       ...,
       [0.985  , 0.784  , 0.1171 , ..., 0.5947 , 0.967  , 0.318  ],
       [0.11755, 0.6265 , 0.1833 , ..., 0.00925, 0.55   , 0.548  ],
       [0.1578 , 0.945  , 0.2764 , ..., 0.7207 , 0.2021 , 0.4163 ]],

      [[0.1431 , 0.978  , 0.25   , ..., 0.2252 , 0.1288 , 0.3447 ],
       [0.3142 , 0.6055 , 0.2864 , ..., 0.9106 , 0.1136 , 0.272  ],
       [0.6147 , 0.847  , 0.396  , ..., 0.554  , 0.01836, 0.1515 ],
       ...,
       [0.2874 , 0.9575 , 0.9766 , ..., 0.2634 , 0.9106 , 0.3525 ],
       [0.978  , 0.71   , 0.4148 , ..., 0.4512 , 0.576  , 0.05423],
       [0.4338 , 0.08575, 0.784  , ..., 0.676  , 0.933  , 0.494  ]],

      [[0.1296 , 0.01025, 0.8755 , ..., 0.1215 , 0.1252 , 0.281  ],
       [0.7437 , 0.05063, 0.8096 , ..., 0.8633 , 0.8735 , 0.5415 ],
       [0.6567 , 0.9146 , 0.533  , ..., 0.3657 , 0.9775 , 0.416  ],
       ...,
       [0.1262 , 0.978  , 0.2334 , ..., 0.3152 , 0.2095 , 0.8735 ],
       [0.539  , 0.775  , 0.77   , ..., 0.5337 , 0.308  , 0.617  ],
       [0.8247 , 0.5903 , 0.3923 , ..., 0.78   , 0.9717 , 0.4485 ]]]]

Outputs:
  Y: shape=(2, 9, 4, 8), dtype=float16
    [[[[0.4565, 0.4722, 0.523 , ..., 0.5005, 0.4812, 0.5786],
       [0.458 , 0.4365, 0.5674, ..., 0.5024, 0.5005, 0.6333],
       [0.4712, 0.4397, 0.572 , ..., 0.5312, 0.4714, 0.591 ],
       [0.4604, 0.448 , 0.5513, ..., 0.5303, 0.4814, 0.601 ]],

      [[0.4646, 0.4583, 0.5293, ..., 0.5044, 0.4863, 0.581 ],
       [0.455 , 0.4417, 0.5703, ..., 0.513 , 0.487 , 0.628 ],
       [0.4663, 0.4407, 0.5615, ..., 0.529 , 0.4727, 0.5933],
       [0.479 , 0.4465, 0.554 , ..., 0.5254, 0.4841, 0.581 ]],

      [[0.4517, 0.4746, 0.525 , ..., 0.507 , 0.478 , 0.571 ],
       [0.459 , 0.4436, 0.5767, ..., 0.5186, 0.4827, 0.6274],
       [0.4695, 0.4456, 0.5645, ..., 0.5303, 0.4695, 0.588 ],
       [0.4832, 0.456 , 0.5654, ..., 0.533 , 0.4768, 0.5767]],

      ...,

      [[0.5464, 0.4841, 0.568 , ..., 0.4487, 0.5005, 0.618 ],
       [0.561 , 0.4775, 0.5386, ..., 0.4868, 0.5537, 0.5376],
       [0.5625, 0.4692, 0.5664, ..., 0.4783, 0.517 , 0.585 ],
       [0.5317, 0.4702, 0.552 , ..., 0.487 , 0.5317, 0.5674]],

      [[0.5435, 0.4827, 0.5664, ..., 0.4473, 0.498 , 0.613 ],
       [0.5713, 0.4663, 0.55  , ..., 0.4805, 0.5615, 0.53  ],
       [0.567 , 0.4734, 0.5645, ..., 0.4763, 0.5137, 0.587 ],
       [0.5337, 0.4688, 0.547 , ..., 0.4805, 0.5396, 0.5693]],

      [[0.553 , 0.4917, 0.5664, ..., 0.4568, 0.5054, 0.613 ],
       [0.567 , 0.4707, 0.5464, ..., 0.4817, 0.5566, 0.54  ],
       [0.5703, 0.466 , 0.567 , ..., 0.4705, 0.5186, 0.5737],
       [0.5293, 0.471 , 0.543 , ..., 0.4805, 0.5312, 0.5728]]],


     [[[0.46  , 0.603 , 0.424 , ..., 0.4155, 0.527 , 0.458 ],
       [0.483 , 0.6294, 0.4792, ..., 0.416 , 0.6055, 0.444 ],
       [0.4775, 0.6045, 0.4565, ..., 0.4143, 0.5664, 0.4448],
       [0.4692, 0.6245, 0.4358, ..., 0.4404, 0.598 , 0.4285]],

      [[0.4578, 0.599 , 0.4304, ..., 0.4202, 0.53  , 0.4531],
       [0.4937, 0.62  , 0.4722, ..., 0.424 , 0.611 , 0.4446],
       [0.475 , 0.6143, 0.4624, ..., 0.4182, 0.5596, 0.4385],
       [0.4631, 0.626 , 0.4348, ..., 0.4346, 0.5884, 0.4365]],

      [[0.465 , 0.6035, 0.4255, ..., 0.4248, 0.5234, 0.4507],
       [0.4963, 0.626 , 0.4705, ..., 0.4219, 0.6035, 0.444 ],
       [0.4622, 0.615 , 0.4634, ..., 0.4053, 0.5674, 0.4463],
       [0.453 , 0.623 , 0.4297, ..., 0.4285, 0.5957, 0.4434]],

      ...,

      [[0.5327, 0.534 , 0.5625, ..., 0.55  , 0.5537, 0.4788],
       [0.4932, 0.507 , 0.556 , ..., 0.5186, 0.5034, 0.519 ],
       [0.5234, 0.5312, 0.5586, ..., 0.562 , 0.5537, 0.493 ],
       [0.4788, 0.4834, 0.611 , ..., 0.4995, 0.466 , 0.526 ]],

      [[0.5376, 0.542 , 0.5605, ..., 0.548 , 0.5537, 0.4817],
       [0.4856, 0.4968, 0.557 , ..., 0.51  , 0.4897, 0.5293],
       [0.525 , 0.521 , 0.55  , ..., 0.555 , 0.546 , 0.494 ],
       [0.4934, 0.511 , 0.613 , ..., 0.514 , 0.4817, 0.5186]],

      [[0.5312, 0.538 , 0.5625, ..., 0.55  , 0.5547, 0.4792],
       [0.4773, 0.4963, 0.5547, ..., 0.5195, 0.4883, 0.5244],
       [0.5156, 0.548 , 0.5483, ..., 0.5635, 0.547 , 0.506 ],
       [0.51  , 0.5337, 0.582 , ..., 0.5166, 0.4783, 0.5244]]]]
  present_key: shape=(2, 3, 18, 8), dtype=float16
    [[[[5.3027e-01, 1.5527e-01, 9.1016e-01, ..., 3.0933e-01, 8.9941e-01,
        2.1765e-01],
       [3.3838e-01, 6.6211e-01, 6.6016e-01, ..., 5.7324e-01, 7.9529e-02,
        9.2822e-01],
       [5.3418e-01, 2.9688e-01, 1.8799e-01, ..., 6.8652e-01, 2.9761e-01,
        3.2440e-02],
       ...,
       [6.8115e-01, 6.8457e-01, 8.6914e-01, ..., 5.5811e-01, 8.3301e-01,
        8.6084e-01],
       [3.6938e-01, 7.8760e-01, 8.8867e-01, ..., 5.0098e-01, 2.8125e-01,
        2.8174e-01],
       [2.3267e-01, 2.2388e-01, 9.4727e-01, ..., 1.4807e-01, 9.7021e-01,
        9.5312e-01]],

      [[7.7490e-01, 3.5919e-02, 1.1002e-02, ..., 1.3878e-02, 5.4004e-01,
        3.0029e-01],
       [7.7686e-01, 5.9668e-01, 3.6450e-01, ..., 6.9775e-01, 5.8447e-01,
        5.2979e-01],
       [3.7549e-01, 1.6079e-03, 8.1445e-01, ..., 6.6650e-01, 3.2886e-01,
        7.6025e-01],
       ...,
       [3.9355e-01, 7.1338e-01, 9.7559e-01, ..., 6.0107e-01, 8.2617e-01,
        1.1041e-01],
       [8.1445e-01, 9.1895e-01, 8.4229e-01, ..., 1.1676e-01, 4.8340e-01,
        2.0715e-01],
       [8.5449e-01, 7.3340e-01, 8.3984e-01, ..., 6.5967e-01, 8.4045e-02,
        1.8262e-01]],

      [[3.9038e-01, 1.1890e-01, 6.1816e-01, ..., 7.5293e-01, 5.0293e-01,
        9.5068e-01],
       [4.6924e-01, 2.0972e-01, 7.7051e-01, ..., 9.5801e-01, 7.3193e-01,
        8.6328e-01],
       [6.0400e-01, 1.1182e-01, 5.8789e-01, ..., 1.3391e-01, 5.7324e-01,
        6.3428e-01],
       ...,
       [5.4199e-01, 3.7646e-01, 9.2871e-01, ..., 8.0371e-01, 9.8291e-01,
        7.0898e-01],
       [6.4258e-01, 9.9902e-01, 7.2949e-01, ..., 9.8779e-01, 9.3750e-02,
        6.5576e-01],
       [1.1945e-01, 5.6934e-01, 5.6738e-01, ..., 5.8447e-01, 6.6260e-01,
        6.7139e-01]]],


     [[[4.7119e-01, 6.2305e-01, 5.7471e-01, ..., 2.4500e-01, 1.3184e-01,
        6.3721e-01],
       [5.0977e-01, 7.2070e-01, 7.2168e-01, ..., 7.3059e-02, 8.6572e-01,
        3.9673e-01],
       [3.1299e-01, 8.4619e-01, 3.5181e-01, ..., 8.0762e-01, 3.7451e-01,
        1.2250e-01],
       ...,
       [4.9683e-01, 4.9902e-01, 4.4312e-01, ..., 6.7725e-01, 9.5557e-01,
        1.3770e-01],
       [9.9731e-02, 2.2778e-01, 1.4148e-01, ..., 9.8050e-05, 6.3965e-01,
        4.8267e-01],
       [1.4832e-01, 3.0591e-01, 5.9424e-01, ..., 1.8433e-01, 1.0815e-01,
        6.4990e-01]],

      [[1.1072e-01, 7.6709e-01, 7.2510e-01, ..., 2.0898e-01, 7.7100e-01,
        8.2861e-01],
       [9.9170e-01, 3.4790e-01, 6.1890e-02, ..., 7.6172e-01, 3.3032e-01,
        5.9961e-01],
       [5.9668e-01, 4.2383e-01, 9.5020e-01, ..., 4.0625e-01, 7.6477e-02,
        6.9385e-01],
       ...,
       [8.7988e-01, 1.5479e-01, 4.8364e-01, ..., 9.8340e-01, 6.4746e-01,
        1.2622e-01],
       [6.8213e-01, 5.4199e-02, 9.1748e-01, ..., 2.9102e-01, 8.9795e-01,
        7.9346e-01],
       [3.1226e-01, 7.6709e-01, 6.3574e-01, ..., 9.7070e-01, 7.3120e-02,
        5.9863e-01]],

      [[4.0283e-01, 7.6953e-01, 9.1309e-01, ..., 3.6377e-01, 7.6367e-01,
        9.9463e-01],
       [9.2316e-03, 3.1055e-01, 6.3477e-01, ..., 7.4268e-01, 4.6899e-01,
        9.4043e-01],
       [5.4834e-01, 6.2988e-02, 2.2607e-01, ..., 6.4893e-01, 2.2797e-02,
        8.1738e-01],
       ...,
       [5.6250e-01, 6.6797e-01, 2.6514e-01, ..., 7.3389e-01, 2.0972e-01,
        6.5527e-01],
       [1.3635e-01, 1.7761e-01, 1.4709e-01, ..., 2.4695e-01, 8.9648e-01,
        9.4385e-01],
       [2.3218e-01, 2.6993e-02, 6.5674e-01, ..., 7.0557e-01, 4.5972e-01,
        2.4805e-01]]]]
  present_value: shape=(2, 3, 18, 8), dtype=float16
    [[[[0.608   , 0.9365  , 0.646   , ..., 0.83    , 0.3618  , 0.2002  ],
       [0.1482  , 0.741   , 0.8584  , ..., 0.368   , 0.3079  , 0.9976  ],
       [0.6387  , 0.824   , 0.777   , ..., 0.1534  , 0.9316  , 0.599   ],
       ...,
       [0.3052  , 0.486   , 0.628   , ..., 0.5186  , 0.7974  , 0.564   ],
       [0.291   , 0.605   , 0.704   , ..., 0.442   , 0.25    , 0.422   ],
       [0.736   , 0.3054  , 0.3667  , ..., 0.1913  , 0.8413  , 0.265   ]],

      [[0.1675  , 0.921   , 0.1498  , ..., 0.0378  , 0.908   , 0.7837  ],
       [0.4946  , 0.6074  , 0.3533  , ..., 0.601   , 0.2115  , 0.6333  ],
       [0.9414  , 0.781   , 0.9795  , ..., 0.2142  , 0.6284  , 0.9487  ],
       ...,
       [0.1678  , 0.7637  , 0.7793  , ..., 0.445   , 0.52    , 0.4966  ],
       [0.2012  , 0.451   , 0.6357  , ..., 0.708   , 0.02582 , 0.3772  ],
       [0.208   , 0.91    , 0.05463 , ..., 0.659   , 0.3806  , 0.9727  ]],

      [[0.631   , 0.527   , 0.1058  , ..., 0.7773  , 0.739   , 0.94    ],
       [0.658   , 0.9775  , 0.7773  , ..., 0.373   , 0.5015  , 0.1617  ],
       [0.2258  , 0.3743  , 0.838   , ..., 0.072   , 0.06244 , 0.7617  ],
       ...,
       [0.4094  , 0.1833  , 0.6226  , ..., 0.218   , 0.5874  , 0.383   ],
       [0.6094  , 0.988   , 0.394   , ..., 0.9004  , 0.1876  , 0.82    ],
       [0.1097  , 0.3174  , 0.01627 , ..., 0.696   , 0.6597  , 0.7627  ]]],


     [[[0.4666  , 0.848   , 0.733   , ..., 0.6416  , 0.5728  , 0.01241 ],
       [0.1087  , 0.946   , 0.3486  , ..., 0.2783  , 0.8955  , 0.927   ],
       [0.961   , 0.6987  , 0.5703  , ..., 0.3208  , 0.9526  , 0.381   ],
       ...,
       [0.932   , 0.4756  , 0.6533  , ..., 0.7554  , 0.933   , 0.3604  ],
       [0.2925  , 0.3213  , 0.3499  , ..., 0.312   , 0.544   , 0.2612  ],
       [0.09314 , 0.007397, 0.4978  , ..., 0.969   , 0.3613  , 0.5234  ]],

      [[0.1431  , 0.978   , 0.25    , ..., 0.2252  , 0.1288  , 0.3447  ],
       [0.3142  , 0.6055  , 0.2864  , ..., 0.9106  , 0.1136  , 0.272   ],
       [0.6147  , 0.847   , 0.396   , ..., 0.554   , 0.01836 , 0.1515  ],
       ...,
       [0.888   , 0.621   , 0.618   , ..., 0.2245  , 0.7417  , 0.8154  ],
       [0.278   , 0.8066  , 0.12445 , ..., 0.486   , 0.1539  , 0.4636  ],
       [0.0608  , 0.7104  , 0.995   , ..., 0.398   , 0.42    , 0.3901  ]],

      [[0.1296  , 0.01025 , 0.8755  , ..., 0.1215  , 0.1252  , 0.281   ],
       [0.7437  , 0.05063 , 0.8096  , ..., 0.8633  , 0.8735  , 0.5415  ],
       [0.6567  , 0.9146  , 0.533   , ..., 0.3657  , 0.9775  , 0.416   ],
       ...,
       [0.6797  , 0.85    , 0.5127  , ..., 0.04388 , 0.4456  , 0.3848  ],
       [0.8438  , 0.04922 , 0.589   , ..., 0.2152  , 0.3254  , 0.4387  ],
       [0.165   , 0.4075  , 0.2546  , ..., 0.1492  , 0.01906 , 0.75    ]]]]

test_cc_attention_4d_scaled

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 0.009999999776482582
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.,  0.],
       [ 0.,  1.]],

      [[ 1.,  1.],
       [-1.,  1.]]]]
  V: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.,  2.],
       [ 3.,  4.]],

      [[-1.,  0.],
       [ 0.,  1.]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1.995     ,  2.9950001 ],
       [ 2.0049999 ,  3.0049999 ]],

      [[-0.5025    ,  0.49750003],
       [-0.5049998 ,  0.49500015]]]]

test_cc_attention_4d_softcap

Node:
  Attention(Q, K, V) -> (Y)
  Attributes:
    scale = 1.0
    softcap = 0.5
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.15169363,  0.6033754 ],
       [-0.15169363,  0.755069  ]],

      [[ 0.67558366, -0.4552321 ],
       [ 0.29168454, -0.3030643 ]]]]

test_cc_attention_4d_softcap_neginf_mask

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    softcap = 0.5
Inputs:
  Q: shape=(1, 1, 4, 8), dtype=float32
    [[[[0.7167289 , 0.6484601 , 0.46497524, 0.34117728, 0.55324554, 0.47324938,
        0.8506881 , 0.04282242],
       [0.4230454 , 0.07251757, 0.7479862 , 0.1587475 , 0.6062188 , 0.806189  ,
        0.17963564, 0.45366198],
       [0.13773894, 0.34766454, 0.74359787, 0.6698144 , 0.26058674, 0.5744466 ,
        0.1911546 , 0.04261738],
       [0.12887597, 0.70876426, 0.2550329 , 0.06627256, 0.7609782 , 0.708198  ,
        0.7275301 , 0.50939566]]]]
  K: shape=(1, 1, 6, 8), dtype=float32
    [[[[0.81594235, 0.868624  , 0.75029296, 0.81736696, 0.16860753, 0.88723135,
        0.19616836, 0.57798064],
       [0.65411955, 0.7793319 , 0.9017522 , 0.47066295, 0.6826957 , 0.4715004 ,
        0.6312656 , 0.7740394 ],
       [0.3452173 , 0.23894429, 0.18115634, 0.68568075, 0.17885476, 0.65862894,
        0.77072376, 0.40022814],
       [0.18269259, 0.8124025 , 0.6827437 , 0.437518  , 0.08335429, 0.7401209 ,
        0.15409368, 0.22070706],
       [0.8506275 , 0.14277291, 0.5163776 , 0.90439737, 0.4630888 , 0.33560514,
        0.58655113, 0.4527613 ],
       [0.923729  , 0.45124698, 0.30754632, 0.9676665 , 0.52557784, 0.8928356 ,
        0.6388969 , 0.266801  ]]]]
  V: shape=(1, 1, 6, 8), dtype=float32
    [[[[0.91515577, 0.08878785, 0.03561068, 0.29355663, 0.7839695 , 0.30121332,
        0.5416486 , 0.11313885],
       [0.8851937 , 0.48614627, 0.05551815, 0.7825784 , 0.75917256, 0.1368118 ,
        0.08289552, 0.0944168 ],
       [0.55269563, 0.13022405, 0.6187148 , 0.7015471 , 0.09712279, 0.74281126,
        0.35029292, 0.7578389 ],
       [0.2365092 , 0.9160407 , 0.11045456, 0.80876344, 0.40573037, 0.77204376,
        0.58065724, 0.93201846],
       [0.45860708, 0.3792408 , 0.2441163 , 0.12815326, 0.74603045, 0.76796633,
        0.5889299 , 0.96119374],
       [0.4440869 , 0.19609386, 0.56120396, 0.44926733, 0.1321832 , 0.664661  ,
        0.26817727, 0.20522803]]]]
  attn_mask: shape=(4, 6), dtype=float32
    [[  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf]]

Outputs:
  Y: shape=(1, 1, 4, 8), dtype=float32
    [[[[0.65484625, 0.40108806, 0.20163937, 0.64396936, 0.5180408 , 0.48061964,
        0.38587224, 0.4638726 ],
       [0.6566198 , 0.40271884, 0.19770749, 0.64299566, 0.5223832 , 0.47754967,
        0.38596648, 0.46013588],
       [0.65369874, 0.40479586, 0.19766673, 0.6429548 , 0.5210476 , 0.48034573,
        0.38840714, 0.4638093 ],
       [0.65444815, 0.40250054, 0.20119615, 0.64485633, 0.518264  , 0.48043442,
        0.38539407, 0.46407378]]]]

test_cc_attention_4d_softcap_neginf_mask_poison

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    softcap = 0.5
Inputs:
  Q: shape=(1, 1, 4, 8), dtype=float32
    [[[[0.7167289 , 0.6484601 , 0.46497524, 0.34117728, 0.55324554, 0.47324938,
        0.8506881 , 0.04282242],
       [0.4230454 , 0.07251757, 0.7479862 , 0.1587475 , 0.6062188 , 0.806189  ,
        0.17963564, 0.45366198],
       [0.13773894, 0.34766454, 0.74359787, 0.6698144 , 0.26058674, 0.5744466 ,
        0.1911546 , 0.04261738],
       [0.12887597, 0.70876426, 0.2550329 , 0.06627256, 0.7609782 , 0.708198  ,
        0.7275301 , 0.50939566]]]]
  K: shape=(1, 1, 6, 8), dtype=float32
    [[[[0.81594235, 0.868624  , 0.75029296, 0.81736696, 0.16860753, 0.88723135,
        0.19616836, 0.57798064],
       [0.65411955, 0.7793319 , 0.9017522 , 0.47066295, 0.6826957 , 0.4715004 ,
        0.6312656 , 0.7740394 ],
       [0.3452173 , 0.23894429, 0.18115634, 0.68568075, 0.17885476, 0.65862894,
        0.77072376, 0.40022814],
       [0.18269259, 0.8124025 , 0.6827437 , 0.437518  , 0.08335429, 0.7401209 ,
        0.15409368, 0.22070706],
       [0.8506275 , 0.14277291, 0.5163776 , 0.90439737, 0.4630888 , 0.33560514,
        0.58655113, 0.4527613 ],
       [0.923729  , 0.45124698, 0.30754632, 0.9676665 , 0.52557784, 0.8928356 ,
        0.6388969 , 0.266801  ]]]]
  V: shape=(1, 1, 6, 8), dtype=float32
    [[[[9.15155768e-01, 8.87878537e-02, 3.56106758e-02, 2.93556631e-01,
        7.83969522e-01, 3.01213324e-01, 5.41648626e-01, 1.13138855e-01],
       [8.85193706e-01, 4.86146271e-01, 5.55181503e-02, 7.82578409e-01,
        7.59172559e-01, 1.36811793e-01, 8.28955173e-02, 9.44167972e-02],
       [5.52695632e-01, 1.30224049e-01, 6.18714809e-01, 7.01547086e-01,
        9.71227884e-02, 7.42811263e-01, 3.50292921e-01, 7.57838905e-01],
       [2.36509204e-01, 9.16040719e-01, 1.10454559e-01, 8.08763444e-01,
        4.05730367e-01, 7.72043765e-01, 5.80657244e-01, 9.32018459e-01],
       [1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03,
        1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03],
       [1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03,
        1.00000000e+03, 1.00000000e+03, 1.00000000e+03, 1.00000000e+03]]]]
  attn_mask: shape=(4, 6), dtype=float32
    [[  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf],
     [  0.,   0.,   0.,   0., -inf, -inf]]

Outputs:
  Y: shape=(1, 1, 4, 8), dtype=float32
    [[[[0.65484625, 0.40108806, 0.20163937, 0.64396936, 0.5180408 , 0.48061964,
        0.38587224, 0.4638726 ],
       [0.6566198 , 0.40271884, 0.19770749, 0.64299566, 0.5223832 , 0.47754967,
        0.38596648, 0.46013588],
       [0.65369874, 0.40479586, 0.19766673, 0.6429548 , 0.5210476 , 0.48034573,
        0.38840714, 0.4638093 ],
       [0.65444815, 0.40250054, 0.20119615, 0.64485633, 0.518264  , 0.48043442,
        0.38539407, 0.46407378]]]]

test_cc_attention_4d_with_past_and_present

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value)
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.10764056,  0.45608166],
       [-0.31939295,  0.57814157]],

      [[ 0.48057237, -0.20996696],
       [ 0.05822894,  0.11373253]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul

Node:
  Attention(Q, K, V, "", past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.10764056,  0.45608166],
       [-0.31939295,  0.57814157]],

      [[ 0.48057237, -0.20996696],
       [ 0.05822894,  0.11373253]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.35355338,  0.        ,  0.70710677,  0.35355338,  0.        ],
       [-0.35355338,  0.35355338,  0.        ,  0.35355338,  0.70710677]],

      [[ 0.35355338,  0.17677669,  0.        ,  0.70710677, -0.08838835],
       [ 0.70710677, -1.0606601 , -1.4142135 ,  0.        ,  0.53033006]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul_bias

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    scale = 0.5
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(2, 5), dtype=float32
    [[ 0. , -0.5, -1. ,  0.2,  0. ],
     [ 0.5,  0. , -0.2, -0.1,  0. ]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[-0.07067307,  0.63377535],
       [-0.23752189,  0.5571051 ]],

      [[ 0.30545416,  0.01247976],
       [ 0.09629637,  0.127618  ]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    scale = 0.5
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 5), dtype=float32
    [[[ 0. , -0.5, -1. ,  0.2,  0. ],
      [ 0.5,  0. , -0.2, -0.1,  0. ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[-0.07067307,  0.63377535],
       [-0.23752189,  0.5571051 ]],

      [[ 0.30545416,  0.01247976],
       [ 0.09629637,  0.127618  ]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    scale = 0.5
    is_causal = 1
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 5), dtype=float32
    [[[ 0. , -0.5, -1. ,  0.2,  0. ],
      [ 0.5,  0. , -0.2, -0.1,  0. ]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.2571047 ,  0.2571047 ],
       [ 0.03885095,  0.39657053]],

      [[ 0.46144247, -0.18698354],
       [ 0.29362845,  0.16985056]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    scale = 0.5
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0. , -0.5, -1. ,  0.2,  0. ],
       [ 0.5,  0. , -0.2, -0.1,  0. ]],

      [[ 0.1,  0.2, -0.3,  0. ,  0.4],
       [-0.4,  0. ,  0.3, -0.2,  0.1]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[-0.07067307,  0.63377535],
       [-0.23752189,  0.5571051 ]],

      [[ 0.34838727, -0.15952764],
       [ 0.18545583, -0.10488772]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal

Node:
  Attention(Q, K, V, attn_mask, past_key, past_value) -> (Y, present_key, present_value, qk_matmul_output)
  Attributes:
    scale = 0.5
    is_causal = 1
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0. , -0.5, -1. ,  0.2,  0. ],
       [ 0.5,  0. , -0.2, -0.1,  0. ]],

      [[ 0.1,  0.2, -0.3,  0. ,  0.4],
       [-0.4,  0. ,  0.3, -0.2,  0.1]]]]
  past_key: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5, -0.5],
       [ 0. ,  0.5]],

      [[ 1. ,  0. ],
       [-0.5,  1. ]]]]
  past_value: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.5,  0.5],
       [-1. ,  0. ]],

      [[ 0. ,  0.5],
       [ 0.5, -0.5]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.2571047 ,  0.2571047 ],
       [ 0.03885095,  0.39657053]],

      [[ 0.61334735, -0.4131356 ],
       [ 0.5664716 , -0.16319034]]]]
  present_key: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 , -0.5 ],
       [ 0.  ,  0.5 ],
       [ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[ 1.  ,  0.  ],
       [-0.5 ,  1.  ],
       [-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  present_value: shape=(1, 2, 5, 2), dtype=float32
    [[[[ 0.5 ,  0.5 ],
       [-1.  ,  0.  ],
       [ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 0.  ,  0.5 ],
       [ 0.5 , -0.5 ],
       [ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  qk_matmul_output: shape=(1, 2, 2, 5), dtype=float32
    [[[[ 0.25  ,  0.    ,  0.5   ,  0.25  ,  0.    ],
       [-0.25  ,  0.25  ,  0.    ,  0.25  ,  0.5   ]],

      [[ 0.25  ,  0.125 ,  0.    ,  0.5   , -0.0625],
       [ 0.5   , -0.75  , -1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_qk_matmul

Node:
  Attention(Q, K, V) -> (Y, "", "", qk_matmul_output)
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.23092115,  0.5444725 ],
       [-0.23092115,  0.77539366]],

      [[ 0.6482418 , -0.37858847],
       [ 0.04638294, -0.08028026]]]]
  qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.70710677,  0.35355338,  0.        ],
       [ 0.        ,  0.35355338,  0.70710677]],

      [[ 0.        ,  0.70710677, -0.08838835],
       [-1.4142135 ,  0.        ,  0.53033006]]]]

test_cc_attention_4d_with_qk_matmul_bias

Node:
  Attention(Q, K, V, attn_mask) -> (Y, "", "", qk_matmul_output)
  Attributes:
    scale = 0.5
    qk_matmul_output_mode = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(2, 3), dtype=float32
    [[ 0. , -0.5, -1. ],
     [ 0.5,  0. , -0.2]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.458196  ,  0.41020232],
       [ 0.0697852 ,  0.6150191 ]],

      [[ 0.9921614 , -0.7460807 ],
       [ 0.3994022 , -0.34422377]]]]
  qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.5   ,  0.25  ,  0.    ],
       [ 0.    ,  0.25  ,  0.5   ]],

      [[ 0.    ,  0.5   , -0.0625],
       [-1.    ,  0.    ,  0.375 ]]]]

test_cc_attention_4d_with_qk_matmul_softcap

Node:
  Attention(Q, K, V) -> (Y, "", "", qk_matmul_output)
  Attributes:
    scale = 1.0
    softcap = 0.5
    qk_matmul_output_mode = 2
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.15169363,  0.6033754 ],
       [-0.15169363,  0.755069  ]],

      [[ 0.67558366, -0.4552321 ],
       [ 0.29168454, -0.3030643 ]]]]
  qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
    [[[[ 0.4820138 ,  0.3807971 ,  0.        ],
       [ 0.        ,  0.3807971 ,  0.4820138 ]],

      [[ 0.        ,  0.4820138 , -0.12245933],
       [-0.49966466,  0.        ,  0.45257413]]]]

test_cc_attention_4d_with_qk_matmul_softmax

Node:
  Attention(Q, K, V) -> (Y, "", "", qk_matmul_output)
  Attributes:
    qk_matmul_output_mode = 3
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.23092115,  0.5444725 ],
       [-0.23092115,  0.77539366]],

      [[ 0.6482418 , -0.37858847],
       [ 0.04638294, -0.08028026]]]]
  qk_matmul_output: shape=(1, 2, 2, 3), dtype=float32
    [[[[0.45552748, 0.31986618, 0.22460635],
       [0.22460635, 0.31986618, 0.45552748]],

      [[0.25358054, 0.5142905 , 0.23212898],
       [0.08261943, 0.33983436, 0.5775462 ]]]]

test_cc_attention_causal_boolmask_nan_robustness

Node:
  Attention(Q, K, V, attn_mask) -> (Y)
  Attributes:
    scale = 0.5
    is_causal = 1
Inputs:
  Q: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 1. ,  0. ],
       [ 0. ,  1. ]],

      [[ 0.5,  0.5],
       [ 1. , -1. ]]]]
  K: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.5 ,  0.5 ],
       [ 0.  ,  1.  ]],

      [[-1.  ,  1.  ],
       [ 1.  ,  1.  ],
       [ 0.25, -0.5 ]]]]
  V: shape=(1, 2, 3, 2), dtype=float32
    [[[[ 1.  ,  0.  ],
       [ 0.  ,  1.  ],
       [-1.  ,  1.  ]],

      [[ 2.  , -2.  ],
       [ 0.5 ,  0.25],
       [-0.5 ,  0.  ]]]]
  attn_mask: shape=(1, 2, 2, 3), dtype=bool
    [[[[False, False, False],
       [ True,  True,  True]],

      [[False, False, False],
       [ True,  True,  True]]]]

Outputs:
  Y: shape=(1, 2, 2, 2), dtype=float32
    [[[[ 0.       ,  0.       ],
       [ 0.4378235,  0.5621765]],

      [[ 0.       ,  0.       ],
       [ 0.9034121, -0.3551182]]]]

test_cc_shape_inference_tiny_llm

Node:
  Attention(query, key, value, attn_bias, past_key, past_value) -> (attn_out, present_key, present_value)
  Attributes:
    q_num_heads = 4
    kv_num_heads = 4
    is_causal = 1