import torch


def angle_axis_to_rotation_matrix(angle_axis):
    """Convert batch of 3D angle-axis vectors into a batch of 3D rotation matrices

    Arguments:
        angle_axis: (b, 3) Torch tensor,
            batch of 3D angle-axis vectors

    Return Values:
        rotation_matrix: (b, 3, 3) Torch tensor,
            batch of 3D rotation matrices
    """
    assert angle_axis.shape[-1] == 3, "Angle-axis vector must be a (*, 3) tensor, received {}".format(
        angle_axis.shape)

    def angle_axis_to_rotation_matrix_rodrigues(angle_axis, theta2):
        theta = torch.sqrt(theta2).unsqueeze(-1) # bx1
        r = angle_axis / theta # bx3
        rx = r[..., 0] # b
        ry = r[..., 1] # b
        rz = r[..., 2] # b
        r_skew = torch.zeros_like(r).unsqueeze(-1).repeat_interleave(3, dim=-1) # bx3x3
        r_skew[..., 2, 1] = rx
        r_skew[..., 1, 2] = -rx
        r_skew[..., 0, 2] = ry
        r_skew[..., 2, 0] = -ry
        r_skew[..., 1, 0] = rz
        r_skew[..., 0, 1] = -rz
        R = torch.eye(3, dtype=r.dtype, device=r.device).unsqueeze(0) \
            + theta.sin().unsqueeze(-1) * r_skew \
            + (1.0 - theta.cos().unsqueeze(-1)) * torch.matmul(r_skew, r_skew) # bx3x3
        return R

    def angle_axis_to_rotation_matrix_taylor(angle_axis):
        rx, ry, rz = torch.chunk(angle_axis, 3, dim=-1)
        ones = torch.ones_like(rx)
        R = torch.cat([ones, -rz, ry, rz, ones, -rx, -ry, rx, ones], dim=1).view(-1, 3, 3)
        return R

    theta2 = torch.einsum('bi,bi->b', (angle_axis, angle_axis))

    eps = 1e-6
    if (theta2 > eps).all():
        rotation_matrix = angle_axis_to_rotation_matrix_rodrigues(angle_axis, theta2)
    else:
        rotation_matrix = angle_axis_to_rotation_matrix_taylor(angle_axis)
        rotation_matrix_rodrigues = angle_axis_to_rotation_matrix_rodrigues(angle_axis, theta2)
        # Iterate over batch dimension
        # Note: cannot use masking or torch.where because any NaNs in the gradient
        # of the unused branch get propagated
        # See: https://github.com/pytorch/pytorch/issues/9688
        for b in range(angle_axis.shape[0]):
            if theta2[b, ...] > eps:
                rotation_matrix[b, ...] = rotation_matrix_rodrigues[b:(b+1), ...]
    return rotation_matrix