[DiadkoShmek]

tfp.math.lambertw is smooth at zero, but its masked W / (z * (1 + W)) derivative gives incorrect higher derivatives there. On the JAX substrate, reverse-mode W''(0) is NaN and forward-mode W''(0) is zero, although the correct value is -2. The Hessian at zero of 0.5 * (W(z) - 1)**2 is consequently 1 instead of 3.

Use the equivalent exp(-W) / (1 + W) expression in both derivative rules so differentiation remains smooth through zero. Add float32/float64 regressions for derivatives through fourth order and batched second derivatives containing zeros.

Validation:

• 18 native Lambert W tests pass on the officially generated JAX CPU substrate; the four new parameterized cases fail on the original implementation.
• A separate 70-digit mpmath oracle checks 361 value/derivative/consumer cases across reverse/forward AD, vmap and jit: 61 baseline failures become zero.
• Independent review checks actual representable inputs, coupled Hessians, shapes and nonuniform tangents/cotangents. Both zero-initialized loss Hessians become the analytic result. That review passes 33 of 34 aggregate comparisons: the remaining float32 forward third derivative at approximately 1e12 has 0.143% relative error; the base has the wrong sign. Large-argument higher-order precision remains limited.
• Repository pylint retains the same seven pre-existing superfluous-parens diagnostics; no new diagnostics. git diff --check passes.

Tested with JAX/jaxlib 0.7.0 on CPU. TensorFlow custom-VJP registration, GPU execution and Bazel were not exercised. The original owner files are modified; generated substrate files are used for validation and are not included in this patch.

I directed this work and take responsibility for the contribution and its follow-up. AI assistance was used in producing the change.

CI follow-up

The current Python 3.12 changed-file lint workflow installs unpinned Pylint. Its 4.0.5 check flagged seven pre-existing redundant-parenthesis assignments in this already changed source file. Those seven pairs are removed; the full implementation AST (including type comments) and test bytes are unchanged. Both affected files now pass the unmodified project Pylint command. The same three pre-existing configuration UserWarnings remain; no checks or configuration were disabled. Numerical testing was not repeated for this syntax-only amendment; the original numerical evidence and stated residual limits above are retained.