import torch


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

    Args:
        z: The input to the lambertw function.

    Returns:
        An estimated value of lambertw(z).
    """
    log1pz = torch.log1p(z)
    return log1pz * (1. - torch.log1p(log1pz) / (2. + log1pz))