def attention_layer(
    from_tensor,
    to_tensor,
    attention_mask,
    num_attention_heads,
    size_per_head,
    attention_probs_dropout_prob,
    initializer_range,
    input_cache=None,
):
  """Performs multi-headed attention from `from_tensor` to `to_tensor`."""

  def _project(t, name):
    return dense_layer_3d(
        input_tensor=t,
        num_attention_heads=num_attention_heads,
        size_per_head=size_per_head,
        initializer=modeling.create_initializer(initializer_range),
        activation=None,
        name=name)

  query_layer = _project(from_tensor, "query")
  key_layer = _project(to_tensor, "key")
  value_layer = _project(to_tensor, "value")
  output_cache = TransformerCache(keys=key_layer, values=value_layer)

  if input_cache is not None:
    key_layer = tf.concat([input_cache.keys, key_layer], 1)
    value_layer = tf.concat([input_cache.values, value_layer], 1)

  attention_scores = tf.einsum("BTNH,BFNH->BNFT", key_layer, query_layer)
  attention_scores = tf.multiply(attention_scores,
                                 1.0 / math.sqrt(float(size_per_head)))

  # [B, 1, F, T]
  attention_mask = tf.expand_dims(attention_mask, axis=[1])
  attention_scores += (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0
  attention_probs = tf.nn.softmax(attention_scores)
  attention_probs = modeling.dropout(attention_probs,
                                     attention_probs_dropout_prob)

  # [B, F, N, H]
  context_layer = tf.einsum("BNFT,BTNH->BFNH", attention_probs, value_layer)

  return context_layer, output_cache
Performs multi-headed attention from `from_tensor` to `to_tensor`.
language/capwap/utils/transformer_utils.py