def matrix_to_rot6d(rotmat):
    """
    Convert rotation matrix to 6D rotation representation.
    Args:
        rotmat (B x 3 x 3): Batch of rotation matrices.
    Returns:
        6D Rotations (B x 3 x 2).
    """
    return rotmat.view(-1, 3, 3)[:, :, :2]