code: def axis_angle_to_quaternion(axis_angle: Tensor) -> Tensor:
      r"""Convert an axis angle to a quaternion.
  
      The quaternion vector has components in (w, x, y, z) format.
  
      Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
  
      Args:
          axis_angle: tensor with axis angle in radians.
  
      Return:
          tensor with quaternion.
  
      Shape:
          - Input: :math:`(*, 3)` where `*` means, any number of dimensions
          - Output: :math:`(*, 4)`
  
      Example:
          >>> axis_angle = tensor((0., 1., 0.))
          >>> axis_angle_to_quaternion(axis_angle)
          tensor([0.8776, 0.0000, 0.4794, 0.0000])
      """
      if not torch.is_tensor(axis_angle):
          raise TypeError(
              f"Input type is not a Tensor. Got {type(axis_angle)}"
          )
  
      if not axis_angle.shape[-1] == 3:
          raise ValueError(
              f"Input must be a tensor of shape Nx3 or 3. Got {axis_angle.shape}"
          )
  
      # unpack input and compute conversion
      a0: Tensor = axis_angle[..., 0:1]
      a1: Tensor = axis_angle[..., 1:2]
      a2: Tensor = axis_angle[..., 2:3]
      theta_squared: Tensor = a0 * a0 + a1 * a1 + a2 * a2
  
      theta: Tensor = torch.sqrt(theta_squared)
      half_theta: Tensor = theta * 0.5
  
      mask: Tensor = theta_squared > 0.0
      ones: Tensor = torch.ones_like(half_theta)
  
      k_neg: Tensor = 0.5 * ones
      k_pos: Tensor = sin(half_theta) / theta
      k: Tensor = where(mask, k_pos, k_neg)
      w: Tensor = where(mask, cos(half_theta), ones)
  
      quaternion: Tensor = torch.zeros(
          size=(*axis_angle.shape[:-1], 4),
          dtype=axis_angle.dtype,
          device=axis_angle.device,
      )
      quaternion[..., 1:2] = a0 * k
      quaternion[..., 2:3] = a1 * k
      quaternion[..., 3:4] = a2 * k
      quaternion[..., 0:1] = w
      return quaternion
  
imports: from __future__ import annotations
  from typing import Optional
  import torch
  import torch.nn.functional
  from kornia.constants import pi
  from kornia.core import Tensor, concatenate, cos, pad, sin, stack, tensor, where, zeros_like
  from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_SHAPE
  from kornia.utils import deprecated
  from kornia.utils.helpers import _torch_inverse_cast
label_desc: Use image node_id 21 for calling a global function with example usage: axis_angle_to_quaternion(axis_angle) and returns: quaternion