import torch


def so3_log_abs_det_jacobian(x):
    """
    Return element wise log abs det jacobian of exponential map
    :param x: Algebra tensor of shape (..., 3)
    :return: Tensor of shape (..., 3)

    Removable pole: (2-2 cos x)/x^2 -> 1-x^2/12 as x->0
    """
    x_norm = x.double().norm(dim=-1)
    mask = x_norm > 1e-10
    x_norm = torch.where(mask, x_norm, torch.ones_like(x_norm))

    ratio = torch.where(
        mask, (2 - 2 * torch.cos(x_norm)) / x_norm ** 2, 1 - x_norm ** 2 / 12
    )
    return torch.log(ratio).to(x.dtype)