def angle_axis_to_rotation_matrix(angle_axis: torch.Tensor) -> torch.Tensor:
    'Convert 3d vector of axis-angle rotation to 3x3 rotation matrix\n\n    Args:\n        angle_axis (torch.Tensor): tensor of 3d vector of axis-angle rotations.\n\n    Returns:\n        torch.Tensor: tensor of 3x3 rotation matrices.\n\n    Shape:\n        - Input: :math:`(N, 3)`\n        - Output: :math:`(N, 3, 3)`\n\n    Example:\n        >>> input = torch.rand(1, 3)  # Nx3\n        >>> output = angle_axis_to_rotation_matrix(input)  # Nx3x3\n    '
    if (not isinstance(angle_axis, torch.Tensor)):
        raise TypeError('Input type is not a torch.Tensor. Got {}'.format(type(angle_axis)))
    if (not (angle_axis.shape[(- 1)] == 3)):
        raise ValueError('Input size must be a (*, 3) tensor. Got {}'.format(angle_axis.shape))

    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-06):
        k_one = 1.0
        theta = torch.sqrt(theta2)
        wxyz = (angle_axis / (theta + eps))
        (wx, wy, wz) = torch.chunk(wxyz, 3, dim=1)
        cos_theta = torch.cos(theta)
        sin_theta = torch.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 = torch.cat([r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
        return rotation_matrix.view((- 1), 3, 3)

    def _compute_rotation_matrix_taylor(angle_axis):
        (rx, ry, rz) = torch.chunk(angle_axis, 3, dim=1)
        k_one = torch.ones_like(rx)
        rotation_matrix = torch.cat([k_one, (- rz), ry, rz, k_one, (- rx), (- ry), rx, k_one], dim=1)
        return rotation_matrix.view((- 1), 3, 3)
    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
    theta2 = torch.squeeze(theta2, dim=1)
    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
    eps = 1e-06
    mask = (theta2 > eps).view((- 1), 1, 1).to(theta2.device)
    mask_pos = mask.type_as(theta2)
    mask_neg = (mask == False).type_as(theta2)
    batch_size = angle_axis.shape[0]
    rotation_matrix = torch.eye(3).to(angle_axis.device).type_as(angle_axis)
    rotation_matrix = rotation_matrix.view(1, 3, 3).repeat(batch_size, 1, 1)
    rotation_matrix[..., :3, :3] = ((mask_pos * rotation_matrix_normal) + (mask_neg * rotation_matrix_taylor))
    return rotation_matrix
Convert 3d vector of axis-angle rotation to 3x3 rotation matrix

Args:
    angle_axis (torch.Tensor): tensor of 3d vector of axis-angle rotations.

Returns:
    torch.Tensor: tensor of 3x3 rotation matrices.

Shape:
    - Input: :math:`(N, 3)`
    - Output: :math:`(N, 3, 3)`

Example:
    >>> input = torch.rand(1, 3)  # Nx3
    >>> output = angle_axis_to_rotation_matrix(input)  # Nx3x3
def angle_axis_to_rotation_matrix(angle_axis: torch.Tensor) -> torch.Tensor:
    'Convert 3d vector of axis-angle rotation to 3x3 rotation matrix\n\n    Args:\n        angle_axis (torch.Tensor): tensor of 3d vector of axis-angle rotations.\n\n    Returns:\n        torch.Tensor: tensor of 3x3 rotation matrices.\n\n    Shape:\n        - Input: :math:`(N, 3)`\n        - Output: :math:`(N, 3, 3)`\n\n    Example:\n        >>> input = torch.rand(1, 3)  # Nx3\n        >>> output = angle_axis_to_rotation_matrix(input)  # Nx3x3\n    '
    if (not isinstance(angle_axis, torch.Tensor)):
        raise TypeError('Input type is not a torch.Tensor. Got {}'.format(type(angle_axis)))
    if (not (angle_axis.shape[(- 1)] == 3)):
        raise ValueError('Input size must be a (*, 3) tensor. Got {}'.format(angle_axis.shape))

    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-06):
        k_one = 1.0
        theta = torch.sqrt(theta2)
        wxyz = (angle_axis / (theta + eps))
        (wx, wy, wz) = torch.chunk(wxyz, 3, dim=1)
        cos_theta = torch.cos(theta)
        sin_theta = torch.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 = torch.cat([r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
        return rotation_matrix.view((- 1), 3, 3)

    def _compute_rotation_matrix_taylor(angle_axis):
        (rx, ry, rz) = torch.chunk(angle_axis, 3, dim=1)
        k_one = torch.ones_like(rx)
        rotation_matrix = torch.cat([k_one, (- rz), ry, rz, k_one, (- rx), (- ry), rx, k_one], dim=1)
        return rotation_matrix.view((- 1), 3, 3)
    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
    theta2 = torch.squeeze(theta2, dim=1)
    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
    eps = 1e-06
    mask = (theta2 > eps).view((- 1), 1, 1).to(theta2.device)
    mask_pos = mask.type_as(theta2)
    mask_neg = (mask == False).type_as(theta2)
    batch_size = angle_axis.shape[0]
    rotation_matrix = torch.eye(3).to(angle_axis.device).type_as(angle_axis)
    rotation_matrix = rotation_matrix.view(1, 3, 3).repeat(batch_size, 1, 1)
    rotation_matrix[..., :3, :3] = ((mask_pos * rotation_matrix_normal) + (mask_neg * rotation_matrix_taylor))
    return rotation_matrix