def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor:
    'Convert an angle axis to a quaternion.\n    The quaternion vector has components in (x, y, z, w) format.\n\n    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h\n\n    Args:\n        angle_axis (torch.Tensor): tensor with angle axis.\n\n    Return:\n        torch.Tensor: tensor with quaternion.\n\n    Shape:\n        - Input: :math:`(*, 3)` where `*` means, any number of dimensions\n        - Output: :math:`(*, 4)`\n\n    Example:\n        >>> angle_axis = torch.rand(2, 3)  # Nx3\n        >>> quaternion = angle_axis_to_quaternion(angle_axis)  # Nx4\n    '
    if (not torch.is_tensor(angle_axis)):
        raise TypeError('Input type is not a torch.Tensor. Got {}'.format(type(angle_axis)))
    if (not (angle_axis.shape[(- 1)] == 3)):
        raise ValueError('Input must be a tensor of shape Nx3 or 3. Got {}'.format(angle_axis.shape))
    a0: torch.Tensor = angle_axis[..., 0:1]
    a1: torch.Tensor = angle_axis[..., 1:2]
    a2: torch.Tensor = angle_axis[..., 2:3]
    theta_squared: torch.Tensor = (((a0 * a0) + (a1 * a1)) + (a2 * a2))
    theta: torch.Tensor = torch.sqrt(theta_squared)
    half_theta: torch.Tensor = (theta * 0.5)
    mask: torch.Tensor = (theta_squared > 0.0)
    ones: torch.Tensor = torch.ones_like(half_theta)
    k_neg: torch.Tensor = (0.5 * ones)
    k_pos: torch.Tensor = (torch.sin(half_theta) / theta)
    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)
    quaternion: torch.Tensor = torch.zeros_like(angle_axis)
    quaternion[..., 0:1] += (a0 * k)
    quaternion[..., 1:2] += (a1 * k)
    quaternion[..., 2:3] += (a2 * k)
    return torch.cat([w, quaternion], dim=(- 1))
Convert an angle axis to a quaternion.
The quaternion vector has components in (x, y, z, w) format.

Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

Args:
    angle_axis (torch.Tensor): tensor with angle axis.

Return:
    torch.Tensor: tensor with quaternion.

Shape:
    - Input: :math:`(*, 3)` where `*` means, any number of dimensions
    - Output: :math:`(*, 4)`

Example:
    >>> angle_axis = torch.rand(2, 3)  # Nx3
    >>> quaternion = angle_axis_to_quaternion(angle_axis)  # Nx4
def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor:
    'Convert an angle axis to a quaternion.\n    The quaternion vector has components in (x, y, z, w) format.\n\n    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h\n\n    Args:\n        angle_axis (torch.Tensor): tensor with angle axis.\n\n    Return:\n        torch.Tensor: tensor with quaternion.\n\n    Shape:\n        - Input: :math:`(*, 3)` where `*` means, any number of dimensions\n        - Output: :math:`(*, 4)`\n\n    Example:\n        >>> angle_axis = torch.rand(2, 3)  # Nx3\n        >>> quaternion = angle_axis_to_quaternion(angle_axis)  # Nx4\n    '
    if (not torch.is_tensor(angle_axis)):
        raise TypeError('Input type is not a torch.Tensor. Got {}'.format(type(angle_axis)))
    if (not (angle_axis.shape[(- 1)] == 3)):
        raise ValueError('Input must be a tensor of shape Nx3 or 3. Got {}'.format(angle_axis.shape))
    a0: torch.Tensor = angle_axis[..., 0:1]
    a1: torch.Tensor = angle_axis[..., 1:2]
    a2: torch.Tensor = angle_axis[..., 2:3]
    theta_squared: torch.Tensor = (((a0 * a0) + (a1 * a1)) + (a2 * a2))
    theta: torch.Tensor = torch.sqrt(theta_squared)
    half_theta: torch.Tensor = (theta * 0.5)
    mask: torch.Tensor = (theta_squared > 0.0)
    ones: torch.Tensor = torch.ones_like(half_theta)
    k_neg: torch.Tensor = (0.5 * ones)
    k_pos: torch.Tensor = (torch.sin(half_theta) / theta)
    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)
    quaternion: torch.Tensor = torch.zeros_like(angle_axis)
    quaternion[..., 0:1] += (a0 * k)
    quaternion[..., 1:2] += (a1 * k)
    quaternion[..., 2:3] += (a2 * k)
    return torch.cat([w, quaternion], dim=(- 1))