import torch


def rotmat_to_rot6d(R, stack_columns=False):
    """
    :param R: (B, 3, 3)
    :param stack_columns:
        if True, 6D pose representation is [1st col of R, 2nd col of R]^T = [R11, R21, R31, R12, R22, R32]^T
        if False, 6D pose representation is [R11, R12, R21, R22, R31, R32]^T
        Set to False if doing inverse of rot6d_to_rotmat
    :return: rot6d: (B, 6)
    """
    if stack_columns:
        rot6d = torch.cat([R[:, :, 0], R[:, :, 1]], dim=1)
    else:
        rot6d = R[:, :, :2].contiguous().view(-1, 6)
    return rot6d