Looking at this tutorial
matmul_qk = tf.matmul(q, k, transpose_b=True) # (..., seq_len_q, seq_len_k)
is the output of two tensors multiplied of shape:
q: query shape == (Batch, heads, seq_len_q, depth)
k: key shape == (Batch, heads, seq_len_k, depth)
I am failing to understand how the shapes result in this. The second tensor is transposed, so it should be (depth, seq_len_k, heads, Batch)
. I don't see how (Batch, heads, seq_len_q, depth) * (depth, seq_len_k, heads, Batch) -> (Batch, heads, seq_len_q, seq_len_k)
.