pyhgf.math.lambert_w0

pyhgf.math.lambert_w0 = <jax._src.custom_derivatives.custom_jvp object>

Principal branch of the Lambert W function for z >= 0.

Solves w * exp(w) = z via 6 Halley iterations, which yields machine precision for all z >= 0.

A custom JVP is attached so that reverse-mode autodiff does not backprop through the Halley iterations. The implicit relation w * exp(w) = z yields the closed-form derivative

\[\frac{\partial W}{\partial z} = \frac{1}{e^{w}\,(1 + w)},\]

which is well-defined for all z >= 0 (in particular at z = 0, where w = 0 gives ∂W/∂z = 1). This avoids the NaN gradients that would otherwise arise when the Halley iterations’ intermediate exp(w) values overflow.

Parameters:

z (Array | ndarray | bool | number | bool | int | float | complex) – The argument(s) of the Lambert W function, expected to be >= 0.

Returns:

The principal branch \(W_0(z)\) solving w * exp(w) = z.

Return type:

w