import torch


def qrotq(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]

    qvec = q[:, 1:]
    uv = torch.cross(qvec, pv, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)

    pv=(pv + 2 * (q[:, :1] * uv + uuv))

#    return (pv + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
    return torch.cat((pw.unsqueeze(-1), pv), dim=1).view(original_shape)