import torch


def rotation_matrix_from_vectors_torch(vec1, vec2):
    """ `rotation_matrix_from_vectors_rodrigues` but for torch """

    dev = vec1.device
    a, b = (
        vec1 / torch.norm(vec1),
        vec2 / torch.norm(vec2)
    )

    v = torch.cross(a, b)
    c = torch.dot(a, b)
    s = torch.norm(v)
    kmat = torch.tensor([
        [0, -v[2], v[1]],
        [v[2], 0, -v[0]],
        [-v[1], v[0], 0]
    ], requires_grad=True).to(dev)

    return torch.eye(3).to(dev) +\
        kmat +\
        torch.mm(kmat, kmat) * ((1 - c) / (torch.square(s)))