import torch


def quaternion_to_rotmat_jac(q):
    """
    Converts batched quaternions q of shape (batch, 4) to the jacobian of the
    corresponding rotation matrix w.r.t. q of shape (batch, 9, 4)
    """
    qr = q[:, 0:1]
    qi = q[:, 1:2]
    qj = q[:, 2:3]
    qk = q[:, 3:4]
    z = torch.zeros_like(qk)

    r1 = 2. * torch.cat((z, z, -2. * qj, -2. * qk), dim=1)
    r2 = 2. * torch.cat((-qk, qj, qi, -qr), dim=1)
    r3 = 2. * torch.cat((qj, qk, qr, qi), dim=1)
    r4 = 2. * torch.cat((qk, qj, qi, qr), dim=1)
    r5 = 2. * torch.cat((z, -2 * qi, z, -2 * qk), dim=1)
    r6 = 2. * torch.cat((-qi, -qr, qk, qj), dim=1)
    r7 = 2. * torch.cat((-qj, qk, -qr, qi), dim=1)
    r8 = 2. * torch.cat((qi, qr, qk, qj), dim=1)
    r9 = 2. * torch.cat((z, -2 * qi, -2 * qj, z), dim=1)

    return torch.cat((r1.unsqueeze(1), r2.unsqueeze(1), r3.unsqueeze(1),
                      r4.unsqueeze(1), r5.unsqueeze(1), r6.unsqueeze(1),
                      r7.unsqueeze(1), r8.unsqueeze(1), r9.unsqueeze(1)), dim=1)