shape of dist_matrix in MultiSelfAttention

In MultiSelfAttention class:

line 235: dist_matrix = tf.einsum("bwhd,bkhd->bwkh", queries, keys) # dots of (batch, word, key, head)

line 255: select_probs = tf.nn.softmax(dist_matrix) # for each (batch, word, head) probability over keys

The shapes of dist_matrix in the two lines seem not consistent. Should line 235 be changed to:
dist_matrix = tf.einsum("bwhd,bkhd->bwhk", queries, keys) # dots of (batch, word, head, key) ?