import torch
import math


def _taylor_approx(z: torch.Tensor) -> torch.Tensor:
    """Compute an approximation of the lambertw function at z.

    Based on the polynomial expansion in https://arxiv.org/pdf/1003.1628.pdf. An empirical comparison of this polynomial
    expansion against the winitzki approximation found that this one is better when z < -0.2649.

    Args:
        z: The input to the lambertw function.

    Returns:
        An estimated value of lambertw(z).
    """
    p2 = 2 * (1. + math.e * z)
    p = torch.sqrt(p2)
    return -1. + p - p2 / 3. + 0.1527777777 * p2 * p