import torch


def qrotq3(q, p):
    """
    Rotate quaternion(s) p about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 4) for p,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 4).
    """
    assert q.shape[-1] == 4
    assert p.shape[-1] == 4
    assert q.shape[:-1] == p.shape[:-1]

    original_shape = list(p.shape)
    q = q.view(-1, 4)
    p = p.view(-1, 4)
    pw=p[:,0]
    pv=p[:,1:4]

    # Compute outer product
    terms = torch.bmm(q.view(-1, 4, 1), q.view(-1, 1, 4))
    b2=terms[:,1,1]
    c2=terms[:,2,2]
    d2=terms[:,3,3]
    ab=terms[:,0,1]
    ac=terms[:,0,2]
    ad=terms[:,0,3]
    bc=terms[:,1,2]
    bd=terms[:,1,3]
    cd=terms[:,2,3]


    qvec_x=[ 1-2*c2-2*d2, 2*bc-2*ad, 2*ac+2*bd]
    qvec_y=[ 2*bc+2*ad, 1-2*b2-2*d2, 2*cd-2*ab]
    qvec_z=[ 2*bd-2*ac, 2*ab+2*cd, 1-2*b2-2*c2]
    qvec=torch.stack((torch.stack(qvec_x, dim=1), torch.stack(qvec_y, dim=1), torch.stack(qvec_z, dim=1)), dim=1)

    pv=torch.bmm(qvec, pv.unsqueeze(-1)).squeeze()

    return torch.cat((pw.unsqueeze(-1), pv), dim=1).view(original_shape)