import torch


def quat_to_d6(quats:torch.Tensor) -> torch.Tensor: # take (...,4) --> (...,6)
    """This code is adapted from https://github.com/facebookresearch/pytorch3d/blob/master/pytorch3d/transforms/rotation_conversions.py"""
    r, i, j, k = torch.unbind(quats, -1)
    two_s = 2.0 / (quats * quats).sum(-1)
    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    matrix = o.reshape(quats.shape[:-1] + (3, 3))
    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)