code: def multihead_encdec_attention_incremental(query_antecedent,
                                             wq, wo, k, v,
                                             mask,
                                             name="multihead_attention"):
    
    heads, _, kv_channels = k.shape.dims[-3:]
    query_dims = query_antecedent.shape.dims[:-1]
    with tf.variable_scope(name, default_name="multihead_attention"):
      q = mtf.einsum(
          [query_antecedent, wq],
          mtf.Shape(query_dims + [heads, kv_channels]))
      o = dot_product_attention(q, k, v, mask)
      return mtf.einsum([o, wo], query_antecedent.shape)
docs: Incremental attention over encoder (one decode step).
  
    In order to use only one variable containing the four weight matrices
    packed together, we insist that the query and memory antecedents have the
    same dimensionality (io_channels) and that the keys and values have the
    same dimensionality (kv_channels).
  
    memory_dims is a subset of query_dims
  
    Args:
      query_antecedent: a mtf.Tensor with shape query_dims + [io_channels]
      wq: a mtf.Tensor with shape [heads, io_channels, kv_channels]
      wo: a mtf.Tensor with shape [heads, io_channels, kv_channels]
      k: memory_dims + [heads, memory_length, kv_channels]
      v: memory_dims + [heads, memory_length, kv_channels]
      mask: mask Tensor (see attention_mask())
      name: an optional string.
  
    Returns:
      A mtf.Tensor with shape [batch, qlen, io_channels]