def METHOD_NAME(self, x):
    """Reshape and transpose to compute scores.
    Args:
        x: (batch, time, size = n_heads * attn_dim)
    Returns:
        (batch, n_heads, time, attn_dim)
    """
    new_x_shape = x.shape[:-1] + (
        self.num_attention_heads,
        self.attention_head_size,
    )
    return x.reshape(*new_x_shape).transpose(1, 2)