code: def dot_product_attention(q,
                            k,
                            v,
                            mask,
                            dropout=0.0,
                            dropout_broadcast_dims=None,
                            extra_logit=None):
    
    length_kv = k.shape.dims[-2]
    logits_shape = mtf.Shape(q.shape.dims[:-1] + [length_kv])
    logits = mtf.einsum([q, k], logits_shape)
    if mask is not None:
      logits += mask
    weights = mtf.softmax(logits, length_kv, extra_logit=extra_logit)
    if dropout != 0.0:
      weights = mtf.dropout(
          weights, 1.0 - dropout,
          noise_shape=weights.shape - dropout_broadcast_dims)
    depth_v = v.shape.dims[-1]
    outputs_shape = mtf.Shape(q.shape.dims[:-1] + [depth_v])
    outputs = mtf.einsum([weights, v], outputs_shape)
    return outputs
docs: Dot-product attention.
  
    Args:
      q: Tensor with shape [...., length_q, depth_k]. Typically leading dimensions
        are [batch, heads].
      k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
        match with q.
      v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
        match with q.
      mask: mask Tensor (see attention_mask())
      dropout: a float.
      dropout_broadcast_dims: an optional list of mtf.Dimension
      extra_logit: an optional scalar or tensor
  
    Returns:
      Tensor with shape [..., length_q, depth_v].