code: def axis_angle_to_rotation_matrix(axis_angle: Tensor) -> Tensor:
      r"""Convert 3d vector of axis-angle rotation to 3x3 rotation matrix.
  
      Args:
          axis_angle: tensor of 3d vector of axis-angle rotations in radians with shape :math:`(N, 3)`.
  
      Returns:
          tensor of rotation matrices of shape :math:`(N, 3, 3)`.
  
      Example:
          >>> input = tensor([[0., 0., 0.]])
          >>> axis_angle_to_rotation_matrix(input)
          tensor([[[1., 0., 0.],
                   [0., 1., 0.],
                   [0., 0., 1.]]])
  
          >>> input = tensor([[1.5708, 0., 0.]])
          >>> axis_angle_to_rotation_matrix(input)
          tensor([[[ 1.0000e+00,  0.0000e+00,  0.0000e+00],
                   [ 0.0000e+00, -3.6200e-06, -1.0000e+00],
                   [ 0.0000e+00,  1.0000e+00, -3.6200e-06]]])
      """
      if not isinstance(axis_angle, Tensor):
          raise TypeError(
              f"Input type is not a Tensor. Got {type(axis_angle)}"
          )
  
      if not axis_angle.shape[-1] == 3:
          raise ValueError(
              f"Input size must be a (*, 3) tensor. Got {axis_angle.shape}"
          )
  
      def _compute_rotation_matrix(
          axis_angle: Tensor, theta2: Tensor, eps: float = 1e-6
      ) -> Tensor:
          # We want to be careful to only evaluate the square root if the
          # norm of the axis_angle vector is greater than zero. Otherwise
          # we get a division by zero.
          k_one = 1.0
          theta = torch.sqrt(theta2)
          wxyz = axis_angle / (theta + eps)
          wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
          cos_theta = cos(theta)
          sin_theta = sin(theta)
  
          r00 = cos_theta + wx * wx * (k_one - cos_theta)
          r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
          r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
          r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
          r11 = cos_theta + wy * wy * (k_one - cos_theta)
          r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
          r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
          r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
          r22 = cos_theta + wz * wz * (k_one - cos_theta)
          rotation_matrix = concatenate(
              [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1
          )
          return rotation_matrix.view(-1, 3, 3)
  
      def _compute_rotation_matrix_taylor(axis_angle: Tensor) -> Tensor:
          rx, ry, rz = torch.chunk(axis_angle, 3, dim=1)
          k_one = torch.ones_like(rx)
          rotation_matrix = concatenate(
              [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1
          )
          return rotation_matrix.view(-1, 3, 3)
  
      # stolen from ceres/rotation.h
  
      _axis_angle = torch.unsqueeze(axis_angle, dim=1)
      theta2 = torch.matmul(_axis_angle, _axis_angle.transpose(1, 2))
      theta2 = torch.squeeze(theta2, dim=1)
  
      # compute rotation matrices
      rotation_matrix_normal = _compute_rotation_matrix(
          axis_angle, theta2
      )
      rotation_matrix_taylor = _compute_rotation_matrix_taylor(
          axis_angle
      )
  
      # create mask to handle both cases
      eps = 1e-6
      mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
      mask_pos = (mask).type_as(theta2)
      mask_neg = (~mask).type_as(theta2)
  
      # create output pose matrix with masked values
      rotation_matrix = (
          mask_pos * rotation_matrix_normal
          + mask_neg * rotation_matrix_taylor
      )
      return rotation_matrix
  
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 10 for calling a global function with example usage: axis_angle_to_rotation_matrix(axis_angle) and returns: rotation_matrix, rotation_matrix, rotation_matrix