LinearAttention#
Domain:
ai.onnxSince version: 27
Unified linear attention operator for autoregressive decoding (T=1) and prefill (T>1).
The query, key, value, and (where applicable) decay/beta inputs use 3D packed format [B, T, H*D], where heads are flattened into the last dimension; q_num_heads and kv_num_heads are always required and are used to unpack to 4D internally for computation. The optional past_state and present_state are 4D with shape (B, H_kv, d_k, d_v).
Group-query attention (GQA) is supported: q_num_heads must be a positive multiple of
kv_num_heads. When q_num_heads == kv_num_heads this reduces to multi-headed linear
attention; when q_num_heads > kv_num_heads each KV head (and its recurrent state) is
shared by q_num_heads / kv_num_heads query heads (multi-query attention is the
special case kv_num_heads == 1).
The update_rule attribute selects the recurrence type:
“linear”: S_t = S{t-1} + k_t ⊗ v_t; o_t = scale * q_t^T S_t
“gated”: S_t = exp(g_t) * S{t-1} + k_t ⊗ v_t; o_t = scale * q_t^T S_t
“delta”: S_t = S{t-1} + β_t * k_t ⊗ (v_t - S{t-1}^T k_t); o_t = scale * q_t^T S_t
“gated_delta”: S_t = exp(g_t) * S{t-1} + β_t * k_t ⊗ (v_t - exp(g_t) * S{t-1}^T k_t); o_t = scale * q_t^T S_t
where g_t is the decay (in log-space), β_t is the update rate, and ⊗ denotes outer product.
Semantics: Equivalent to running the recurrent update sequentially for each token, but may be implemented using chunk-parallel algorithms for GPU efficiency.
Inputs
query (T): Query vectors with 3D packed shape (B, T, H_q * d_k). Heads are packed into the last dimension.
key (T): Key vectors with 3D packed shape (B, T, H_kv * d_k). Should be L2-normalized for delta/gated_delta modes.
value (T): Value vectors with 3D packed shape (B, T, H_kv * d_v).
past_state (S): Recurrent state from previous step with shape (B, H_kv, d_k, d_v). Always 4D. If not provided, defaults to zeros.
decay (T): Exponential decay gate in log-space. 3D packed shape: (B, T, H_kv * d_k) for per-key-dimension decay (GLA/RWKV-6), or (B, T, H_kv) for per-head scalar decay (DeltaNet/RetNet). Required for ‘gated’ and ‘gated_delta’ modes.
beta (T): Update rate (sigmoid output). 3D packed shape: (B, T, H_kv) or (B, T, 1). Required for ‘delta’ and ‘gated_delta’ modes.
Outputs
output (T): Attention output with 3D packed shape (B, T, H_q * d_v).
present_state (S): Updated recurrent state with shape (B, H_kv, d_k, d_v). Always 4D.
Attributes
chunk_size (int): Chunk size for the chunk-parallel WY decomposition during prefill (T>1). Tuning hint; does not affect output correctness.
kv_num_heads (int): Number of key/value heads. Always required.
q_num_heads (int): Number of query heads. Always required.
scale (float): Output scaling factor. When 0.0 (default), derives d_k = query.shape[-1] / q_num_heads and uses 1/sqrt(d_k). Set explicitly to override.
update_rule (string): The update rule for the linear attention recurrence. One of: ‘linear’, ‘gated’, ‘delta’, ‘gated_delta’. Default is ‘gated_delta’.
Type Constraints
T: Constrain activation input and output types to float16, bfloat16, or float32 tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).
S: Constrain state types to float16, bfloat16, or float32 tensors. Should be float32 or the same as T for numerical stability on long sequences. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).
Examples#
test_cc_linear_attention_decode_step
Node:
LinearAttention(query, key, value, past_state, decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 1, 4), dtype=float32
[[[1. , 0. , 0.5, 0.5]]]
key: shape=(1, 1, 4), dtype=float32
[[[1. , 0. , 0.5, 0.5]]]
value: shape=(1, 1, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5]]]
past_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.5, -0.5],
[ 0. , 0.5]],
[[ 1. , 0. ],
[-0.5, 1. ]]]]
decay: shape=(1, 1, 4), dtype=float32
[[[-0.1, -0.2, -0.3, -0.4]]]
beta: shape=(1, 1, 2), dtype=float32
[[[0.8, 0.9]]]
Outputs:
output: shape=(1, 1, 4), dtype=float32
[[[ 0.6296671 , 1.0673892 , 0.23798102, -0.02875237]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.89048374, 1.5095164 ],
[ 0. , 0.4093654 ]],
[[ 0.8745451 , -0.375822 ],
[-0.2014331 , 0.29449803]]]]
test_cc_linear_attention_delta
Node:
LinearAttention(query, key, value, "", "", beta) -> (output, present_state)
Attributes:
update_rule = "delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
beta: shape=(1, 2, 2), dtype=float32
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.15909901, -0.15909901],
[ 1.4849242 , 1.9798989 , 0.42426407, -0.42426407]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.8 , 1.6 ],
[ 2.1 , 2.8 ]],
[[ 0.525 , -0.525 ],
[-0.07500002, 0.07500002]]]]
test_cc_linear_attention_explicit_scale
Node:
LinearAttention(query, key, value) -> (output, present_state)
Attributes:
update_rule = "linear"
q_num_heads = 2
kv_num_heads = 2
scale = 2.0
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 2. , 4. , 0.5, -0.5],
[ 6. , 8. , 2. , -2. ]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1. , 2. ],
[ 3. , 4. ]],
[[ 0.75, -0.75],
[-0.25, 0.25]]]]
test_cc_linear_attention_fp16
Node:
LinearAttention(query, key, value, "", decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float16
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float16
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float16
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 4), dtype=float16
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
beta: shape=(1, 2, 2), dtype=float16
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 4), dtype=float16
[[[ 0.5654, 1.131 , 0.159 , -0.159 ],
[ 1.485 , 1.98 , 0.429 , -0.429 ]]]
present_state: shape=(1, 2, 2, 2), dtype=float16
[[[[ 0.7607 , 1.521 ],
[ 2.102 , 2.8 ]],
[[ 0.4922 , -0.4922 ],
[-0.11444, 0.11444]]]]
test_cc_linear_attention_gated
Node:
LinearAttention(query, key, value, "", decay) -> (output, present_state)
Attributes:
update_rule = "gated"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 2), dtype=float32
[[[-0.1 , -0.2 ],
[-0.3 , -0.05]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.17677669, -0.17677669],
[ 2.1213202 , 2.828427 , 0.70710677, -0.70710677]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.7408182 , 1.4816364 ],
[ 3. , 4. ]],
[[ 0.7378074 , -0.7378074 ],
[-0.26219264, 0.26219264]]]]
test_cc_linear_attention_gated_delta
Node:
LinearAttention(query, key, value, "", decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 2), dtype=float32
[[[-0.1 , -0.2 ],
[-0.3 , -0.05]]]
beta: shape=(1, 2, 2), dtype=float32
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.15909901, -0.15909901],
[ 1.4849242 , 1.9798989 , 0.42426407, -0.42426407]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.5926546 , 1.1853092 ],
[ 2.1 , 2.8 ]],
[[ 0.51402664, -0.51402664],
[-0.0859734 , 0.0859734 ]]]]
test_cc_linear_attention_gated_delta_beta_scalar
Node:
LinearAttention(query, key, value, "", decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 4), dtype=float32
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
beta: shape=(1, 2, 1), dtype=float32
[[[0.8],
[0.7]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.14142136, -0.14142136],
[ 1.4849242 , 1.9798989 , 0.49883345, -0.49883345]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.7609836 , 1.5219672 ],
[ 2.1 , 2.8 ]],
[[ 0.5206724 , -0.5206724 ],
[-0.18478464, 0.18478464]]]]
test_cc_linear_attention_gated_delta_gqa
Node:
LinearAttention(query, key, value, "", decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 4
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 8), dtype=float32
[[[ 1. , 0. , 0.5, 0.5, -1. , 1. , 0.2, 0.3],
[ 0. , 1. , 1. , -1. , 0.5, 0.5, -0.5, 0.5]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 4), dtype=float32
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
beta: shape=(1, 2, 2), dtype=float32
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 8), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.28284273, 0.56568545, 0. ,
0. , 0.07954951, -0.07954951],
[ 1.4849242 , 1.9798989 , -0.94682753, -0.90370566, 0.13359852,
-0.13359852, -0.21446955, 0.21446955]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.7609836 , 1.5219672 ],
[ 2.1 , 2.8 ]],
[[ 0.49224257, -0.49224257],
[-0.11436889, 0.11436889]]]]
test_cc_linear_attention_gated_delta_mqa
Node:
LinearAttention(query, key, value, "", decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 4
kv_num_heads = 1
Inputs:
query: shape=(1, 2, 8), dtype=float32
[[[ 1. , 0. , 0.5, 0.5, -1. , 1. , 0.2, 0.3],
[ 0. , 1. , 1. , -1. , 0.5, 0.5, -0.5, 0.5]]]
key: shape=(1, 2, 2), dtype=float32
[[[1., 0.],
[0., 1.]]]
value: shape=(1, 2, 2), dtype=float32
[[[1., 2.],
[3., 4.]]]
decay: shape=(1, 2, 2), dtype=float32
[[[-0.1 , -0.2 ],
[-0.05, -0.1 ]]]
beta: shape=(1, 2, 1), dtype=float32
[[[0.8],
[0.7]]]
Outputs:
output: shape=(1, 2, 8), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.28284273, 0.56568545, -0.56568545,
-1.1313709 , 0.11313709, 0.22627418],
[ 1.4849242 , 1.9798989 , -0.94682753, -0.90370566, 1.0115104 ,
1.5280461 , 0.47341377, 0.45185283]]]
present_state: shape=(1, 1, 2, 2), dtype=float32
[[[[0.7609836, 1.5219672],
[2.1 , 2.8 ]]]]
test_cc_linear_attention_gated_per_head_decay
Node:
LinearAttention(query, key, value, "", decay) -> (output, present_state)
Attributes:
update_rule = "gated"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 2), dtype=float32
[[[-0.1 , -0.2 ],
[-0.3 , -0.05]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.17677669, -0.17677669],
[ 2.1213202 , 2.828427 , 0.70710677, -0.70710677]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.7408182 , 1.4816364 ],
[ 3. , 4. ]],
[[ 0.7378074 , -0.7378074 ],
[-0.26219264, 0.26219264]]]]
test_cc_linear_attention_gated_perdim_decay
Node:
LinearAttention(query, key, value, "", decay) -> (output, present_state)
Attributes:
update_rule = "gated"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
decay: shape=(1, 2, 4), dtype=float32
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.17677669, -0.17677669],
[ 2.1213202 , 2.828427 , 0.7145273 , -0.7145273 ]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.95122945, 1.9024589 ],
[ 3. , 4. ]],
[[ 0.715177 , -0.715177 ],
[-0.2953173 , 0.2953173 ]]]]
test_cc_linear_attention_gqa
Node:
LinearAttention(query, key, value) -> (output, present_state)
Attributes:
update_rule = "linear"
q_num_heads = 4
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 8), dtype=float32
[[[ 1. , 0. , 0.5, 0.5, -1. , 1. , 0.2, 0.3],
[ 0. , 1. , 1. , -1. , 0.5, 0.5, -0.5, 0.5]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
Outputs:
output: shape=(1, 2, 8), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.35355338, 0.70710677, 0. ,
0. , 0.08838835, -0.08838835],
[ 2.1213202 , 2.828427 , -1.4142135 , -1.4142135 , 0.17677669,
-0.17677669, -0.35355338, 0.35355338]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1. , 2. ],
[ 3. , 4. ]],
[[ 0.75, -0.75],
[-0.25, 0.25]]]]
test_cc_linear_attention_linear
Node:
LinearAttention(query, key, value) -> (output, present_state)
Attributes:
update_rule = "linear"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.17677669, -0.17677669],
[ 2.1213202 , 2.828427 , 0.70710677, -0.70710677]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1. , 2. ],
[ 3. , 4. ]],
[[ 0.75, -0.75],
[-0.25, 0.25]]]]
test_cc_linear_attention_linear_t1_no_past
Node:
LinearAttention(query, key, value) -> (output, present_state)
Attributes:
update_rule = "linear"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 1, 4), dtype=float32
[[[1. , 0. , 0.5, 0.5]]]
key: shape=(1, 1, 4), dtype=float32
[[[1. , 0. , 0.5, 0.5]]]
value: shape=(1, 1, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5]]]
Outputs:
output: shape=(1, 1, 4), dtype=float32
[[[ 0.70710677, 1.4142135 , 0.17677669, -0.17677669]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1. , 2. ],
[ 0. , 0. ]],
[[ 0.25, -0.25],
[ 0.25, -0.25]]]]
test_cc_linear_attention_no_past_explicit_zeros
Node:
LinearAttention(query, key, value, past_state, decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
past_state: shape=(1, 2, 2, 2), dtype=float32
[[[[0., 0.],
[0., 0.]],
[[0., 0.],
[0., 0.]]]]
decay: shape=(1, 2, 4), dtype=float32
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
beta: shape=(1, 2, 2), dtype=float32
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.56568545, 1.1313709 , 0.15909901, -0.15909901],
[ 1.4849242 , 1.9798989 , 0.4289391 , -0.4289391 ]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.7609836 , 1.5219672 ],
[ 2.1 , 2.8 ]],
[[ 0.49224257, -0.49224257],
[-0.11436889, 0.11436889]]]]
test_cc_linear_attention_prefill_with_past
Node:
LinearAttention(query, key, value, past_state, decay, beta) -> (output, present_state)
Attributes:
update_rule = "gated_delta"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
past_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.5, -0.5],
[ 0. , 0.5]],
[[ 1. , 0. ],
[-0.5, 1. ]]]]
decay: shape=(1, 2, 4), dtype=float32
[[[-0.1 , -0.2 , -0.3 , -0.4 ],
[-0.05, -0.1 , -0.15, -0.2 ]]]
beta: shape=(1, 2, 2), dtype=float32
[[[0.8, 0.9],
[0.7, 0.6]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 0.6296671 , 1.0673892 , 0.23798102, -0.02875237],
[ 1.4849242 , 2.0584745 , 0.8784764 , -0.7037207 ]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.84705436, 1.4358964 ],
[ 2.1 , 2.9111228 ]],
[[ 0.9150809 , -0.53878486],
[-0.32727236, 0.45642647]]]]
test_cc_linear_attention_with_past_state
Node:
LinearAttention(query, key, value, past_state) -> (output, present_state)
Attributes:
update_rule = "linear"
q_num_heads = 2
kv_num_heads = 2
Inputs:
query: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , 1. , -1. ]]]
key: shape=(1, 2, 4), dtype=float32
[[[ 1. , 0. , 0.5, 0.5],
[ 0. , 1. , -0.5, 0.5]]]
value: shape=(1, 2, 4), dtype=float32
[[[ 1. , 2. , 0.5, -0.5],
[ 3. , 4. , -1. , 1. ]]]
past_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 0.5, -0.5],
[ 0. , 0.5]],
[[ 1. , 0. ],
[-0.5, 1. ]]]]
Outputs:
output: shape=(1, 2, 4), dtype=float32
[[[ 1.0606601 , 1.0606601 , 0.35355338, 0.17677669],
[ 2.1213202 , 3.1819804 , 1.767767 , -1.4142135 ]]]
present_state: shape=(1, 2, 2, 2), dtype=float32
[[[[ 1.5 , 1.5 ],
[ 3. , 4.5 ]],
[[ 1.75, -0.75],
[-0.75, 1.25]]]]