View source on GitHub

Computes Lambert W of z element-wise.

tfp.substrates.jax.math.lambertw(
    z, name=None
)

The Lambert W function is the inverse of z = u * exp(u), i. e., it is the function that satisfies u = W(z) * exp(W(z)). The solution cannot be expressed as a composition of elementary functions and is thus part of the special functions in mathematics. See https://en.wikipedia.org/wiki/Lambert_W_function

In general it is a complex-valued function with multiple branches. The k=0 branch is known as the principal branch of the Lambert W function and is implemented here. See also scipy.special.lambertw.

This code returns only the real part of the image of the Lambert W function.

References

Corless, R.M., Gonnet, G.H., Hare, D.E.G. et al. On the LambertW function. Adv Comput Math 5, 329-359 (1996) doi:10.1007/BF02124750