import torch


def quaternion_into_axis_angle(quaternion, device=torch.device("cpu"), dtype=torch.float32):
    """
Takes an quaternion rotation and converts into axis-angle rotation.
https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles

    :param quaternion: 4-element tensor, containig quaternion (q0,q1,q2,q3).
    :param device: Device to allocate new tensors. Default: torch.device("cpu").
    :param dtype: Data type for new tensors. Default: torch.float32.
    :return: (Axis of rotation (3-element tensor), Simple rotation angle (float or 1-element tensor))
    """
    # Simple rotation angle
    angle = torch.nan_to_num(torch.arccos(quaternion[:, 0])) * 2

    # Avoids recalculating this sin.
    sin_angle_2 = torch.sin(angle / 2)

    # Rotation axis
    normalized_axis = torch.zeros((quaternion.shape[0], 3,), device=device, dtype=dtype)
    normalized_axis[:, 0] = quaternion[:, 1] / sin_angle_2
    normalized_axis[:, 1] = quaternion[:, 2] / sin_angle_2
    normalized_axis[:, 2] = quaternion[:, 3] / sin_angle_2

    return normalized_axis, angle