import torch


def axisangle_between_vec(vec1, vec2, eps=1e-6):
    """
    Calculates the rotation between two 3D vectors from @vec1 to @vec2 in axis angle representation
    See https://www.euclideanspace.com/maths/algebra/vectors/angleBetween/ for formula

    Args:
        vec1 (tensor): (..., 3) where final dim represents (x, y, z) 3D vector
        vec2 (tensor): (..., 3) where final dim represents (x, y, z) 3D vector
        eps (float): Stability value below which small values will be mapped to 0

    Returns:
        tensor: (..., 3) where final dim represents (ax, ay, az) axis-angle rotation from @vec1 to @vec2.
    """
    # type: (Tensor, Tensor, float) -> Tensor
    where_small = torch.abs(torch.sum(vec2 - vec1, dim=-1)) < eps
    # Normalize vectors
    vec1 = vec1 / torch.norm(vec1, dim=-1, keepdim=True)
    vec2 = vec2 / torch.norm(vec2, dim=-1, keepdim=True)
    # Calculate angle and axis
    angle = torch.acos(torch.sum(vec1 * vec2, dim=-1, keepdim=True))
    axis = torch.cross(vec1, vec2)
    axis_norm = torch.norm(axis, dim=-1)
    where_small = where_small | (axis_norm < eps)
    axis[where_small] = 0.0
    axis[~where_small] = axis[~where_small] / axis_norm[~where_small].unsqueeze(-1)
    # Combine and return value
    return angle * axis