def _merge_heads(self, tensor, num_heads, attn_head_size):
    '\n        Merges attn_head_size dim and num_attn_heads dim into hidden_size\n        '
    if (len(tensor.shape) == 5):
        tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
    elif (len(tensor.shape) == 4):
        tensor = tensor.permute(0, 2, 1, 3).contiguous()
    else:
        raise ValueError(f'Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}')
    new_shape = (tensor.size()[:(- 2)] + ((num_heads * attn_head_size),))
    return tensor.view(new_shape)
Merges attn_head_size dim and num_attn_heads dim into hidden_size
def _merge_heads(self, tensor, num_heads, attn_head_size):
    '\n        \n        '
    if (len(tensor.shape) == 5):
        tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
    elif (len(tensor.shape) == 4):
        tensor = tensor.permute(0, 2, 1, 3).contiguous()
    else:
        raise ValueError(f'Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}')
    new_shape = (tensor.size()[:(- 2)] + ((num_heads * attn_head_size),))
    return tensor.view(new_shape)