s='no'>value</span>)) <span class='co'># (batch_size, query_elements, value_depth)</span>

<span class='co'># If `value` is not given then internally `value = key` will be used:</span>
<span class='no'>mha</span> <span class='kw'>=</span> <span class='fu'>layer_multi_head_attention</span>(<span class='kw'>head_size</span><span class='kw'>=</span><span class='fl'>128</span>, <span class='kw'>num_heads</span><span class='kw'>=</span><span class='fl'>128</span>)
<span class='no'>query</span> <span class='kw'>=