docs: def combine_heads(self, x):
      """Combine tensor that has been split.
  
      Args:
        x: A tensor [batch_size, num_heads, length, hidden_size/num_heads]
  
      Returns:
        A tensor with shape [batch_size, length, hidden_size]
      """
      with tf.name_scope("combine_heads"):
        batch_size = tf.shape(x)[0]
        length = tf.shape(x)[2]
        x = tf.transpose(x, [0, 2, 1, 3])  # --> [batch, length, num_heads, depth]
        return tf.reshape(x, [batch_size, length, self.hidden_size])