code: def attention(q,
                k,
                v,
                memory_length_dim,
                key_dim,
                value_dim,
                mask=None,
                dropout_rate=0.0,
                dropout_broadcast_dims=None,
                extra_logit=None):
    
    logits = mtf.einsum([q, k], reduced_dims=[key_dim])
    if mask is not None:
      logits += mask
    weights = mtf.softmax(logits, memory_length_dim, extra_logit=extra_logit)
    if dropout_rate != 0.0:
      weights = mtf.dropout(
          weights, 1.0 - dropout_rate,
          noise_shape=weights.shape - dropout_broadcast_dims)
    outputs_shape = q.shape - key_dim + value_dim
    outputs = mtf.einsum([weights, v], outputs_shape)
    return outputs
docs: Dot-product attention - doesn't use positional dimensions.
  
    key_dim is a Dimension representing the channels in the queries and keys
    value_dim is a Dimension representing the channels in values
    memory_length_dim is a Dimension representing the different key/value pairs.
  
    Dimensions of q: other_query_dims + {key_dim}
    Dimensions of k: other_memory_dims + {memory_length_dim, key_dim}
    Dimensions of v: other_memory_dims + {memory_length_dim, value_dim}
    other_memory_dims is a subset of other_query_dims
  
    Typically, other_query_dims={batch, heads, length}
    Typically, other_memory_dims={batch, heads}
  
    Args:
      q: a Tensor
      k: a Tensor
      v: a Tensor
      memory_length_dim: a Dimension
      key_dim: a Dimension
      value_dim: a Dimension
      mask: mask Tensor (see attention_mask())
      dropout_rate: a float.
      dropout_broadcast_dims: an optional list of mtf.Dimension
      extra_logit: an optional scalar or tensor
  
    Returns:
      Tensor with shape q.shape - key_dim + value_dim