import tensorflow as tf

    w = tf.reshape(w, [hidden_size, num_attention_heads, head_size])
    b = tf.get_variable(
        name="bias",
        shape=[num_attention_heads * head_size],
        initializer=tf.zeros_initializer)
    b = tf.reshape(b, [num_attention_heads, head_size])
    ret = tf.einsum("BFH,HND->BFND", input_tensor, w)
    ret += b
  if activation is not None:
    return activation(ret)
  else:
    return ret
