From 75f57784c460ee4a5cbf969f848a388517921de3 Mon Sep 17 00:00:00 2001 From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> Date: Wed, 4 Nov 2026 17:00:53 +0200 Subject: [PATCH 1/2] Fix Lambert W higher derivatives at zero Signed-off-by: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> --- tensorflow_probability/python/math/special.py | 17 +++++--------- .../python/math/special_test.py | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/tensorflow_probability/python/math/special.py b/tensorflow_probability/python/math/special.py index 9d0c472d6a..6562b2a931 100644 --- a/tensorflow_probability/python/math/special.py +++ b/tensorflow_probability/python/math/special.py @@ -1787,12 +1787,9 @@ def _lambert_bwd(aux, g): """Reverse mode impl for lambert.""" z, = aux wz = _lambert_custom_gradient(z) - # At z = 0 the analytic expressions for the gradient results in a 0/0 - # expression. However, the continuous expansion (l'Hospital rule) gives a - # derivative of 1.0 at z = 0. This case has to be handled separately with - # a where clause. - return g * tf.where( - tf.equal(z, 0.), tf.ones([], wz.dtype), wz / (z * (1. + wz))) + # Using z = wz * exp(wz) avoids a removable singularity at z = 0, + # including when differentiating this expression for higher-order gradients. + return g * tf.math.exp(-wz) / (1. + wz) def _lambert_jvp(primals, tangents): @@ -1801,11 +1798,9 @@ def _lambert_jvp(primals, tangents): dz, = tangents wz = _lambert_custom_gradient(z) - # At z = 0 the analytic expressions for the gradient results in a 0/0 - # expression. However, the continuous expansion (l'Hospital rule) gives a - # derivative of 1.0 at z = 0. This case has to be handled separately with - # a where clause. - pz = tf.where(tf.equal(z, 0.), tf.ones([], wz.dtype), wz / (z * (1. + wz))) + # Using z = wz * exp(wz) avoids a removable singularity at z = 0, + # including when differentiating this expression for higher-order gradients. + pz = tf.math.exp(-wz) / (1. + wz) return wz, pz * dz diff --git a/tensorflow_probability/python/math/special_test.py b/tensorflow_probability/python/math/special_test.py index 675bcf2a44..fa80600cc0 100644 --- a/tensorflow_probability/python/math/special_test.py +++ b/tensorflow_probability/python/math/special_test.py @@ -1356,6 +1356,29 @@ def testLambertWGradient(self, value, expected): _, dy_dx = gradient.value_and_gradient(special.lambertw, x) self.assertAllClose(dy_dx, expected) + @parameterized.parameters(tf.float32, tf.float64) + @test_util.numpy_disable_gradient_test + def testLambertWHigherOrderGradientsAtZero(self, dtype): + # The Taylor series is sum_n (-n)**(n - 1) * z**n / n!. + derivative = special.lambertw + x = tf.constant(0., dtype=dtype) + for expected in [1., -2., 9., -64.]: + previous_derivative = derivative + derivative = lambda z, fn=previous_derivative: ( + gradient.value_and_gradient(fn, z)[1]) + self.assertAllClose(self.evaluate(derivative(x)), expected) + + @parameterized.parameters(tf.float32, tf.float64) + @test_util.numpy_disable_gradient_test + def testLambertWSecondDerivative(self, dtype): + # Parameterize by W(z) to obtain independent reference values. + w = np.array([[-0.5, 0., 0.5], [0., 1., 2.]]) + x = tf.constant(w * np.exp(w), dtype=dtype) + derivative = lambda z: gradient.value_and_gradient(special.lambertw, z)[1] + _, second_derivative = gradient.value_and_gradient(derivative, x) + expected = -np.exp(-2. * w) * (w + 2.) / (w + 1.)**3 + self.assertAllClose(self.evaluate(second_derivative), expected) + def testLogGammaCorrection(self): x = half_cauchy.HalfCauchy( loc=8., scale=10.).sample(int(1e3), test_util.test_seed()) From b4ad18d13bb325d8057364c4c31b2ad948da3811 Mon Sep 17 00:00:00 2001 From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> Date: Thu, 5 Nov 2026 01:27:17 +0200 Subject: [PATCH 2/2] Normalize existing assignments for current changed-file lint --- tensorflow_probability/python/math/special.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tensorflow_probability/python/math/special.py b/tensorflow_probability/python/math/special.py index 6562b2a931..8fa616fc75 100644 --- a/tensorflow_probability/python/math/special.py +++ b/tensorflow_probability/python/math/special.py @@ -99,7 +99,7 @@ def _betainc_naive(a, b, x): dtype_orig = dtype_util.common_dtype([a, b, x], tf.float32) # We promote bfloat16 and float16 to float32 to make this function consistent # with the XLA implementation of betainc. - should_promote_dtype = (dtype_orig in _f16bit_dtypes) + should_promote_dtype = dtype_orig in _f16bit_dtypes dtype = tf.float32 if should_promote_dtype else dtype_orig a, b, x = [tf.convert_to_tensor(z, dtype=dtype_orig) for z in [a, b, x]] @@ -356,7 +356,7 @@ def _betainc_der_power_series(a, b, x, dtype, use_power_series): # When x >= a / (a + b), we must apply the symmetry relation given here: # https://dlmf.nist.gov/8.17.E4 # betainc(a, b, x) = 1 - betainc(b, a, 1 - x) - use_symmetry_relation = (safe_x >= safe_a / (safe_a + safe_b)) + use_symmetry_relation = safe_x >= safe_a / (safe_a + safe_b) safe_a_orig = safe_a safe_a = tf.where(use_symmetry_relation, safe_b, safe_a) safe_b = tf.where(use_symmetry_relation, safe_a_orig, safe_b) @@ -433,7 +433,7 @@ def _betainc_partials(a, b, x): dtype_orig = dtype_util.common_dtype([a, b, x], tf.float32) # We promote bfloat16 and float16 to float32 to make this function consistent # with betainc. - should_promote_dtype = (dtype_orig in _f16bit_dtypes) + should_promote_dtype = dtype_orig in _f16bit_dtypes dtype = tf.float32 if should_promote_dtype else dtype_orig numpy_dtype = dtype_util.as_numpy_dtype(dtype) zero = tf.constant(0., dtype=dtype) @@ -645,7 +645,7 @@ def _betaincinv_computation(a, b, y): dtype_orig = dtype_util.common_dtype([a, b, y], tf.float32) # We promote bfloat16 and float16 to float32 to make this function consistent # with betainc. - should_promote_dtype = (dtype_orig in _f16bit_dtypes) + should_promote_dtype = dtype_orig in _f16bit_dtypes dtype = tf.float32 if should_promote_dtype else dtype_orig numpy_dtype = dtype_util.as_numpy_dtype(dtype) zero = numpy_dtype(0.) @@ -677,7 +677,7 @@ def _betaincinv_computation(a, b, y): use_symmetry_relation = (error_at_half < zero) & a_and_b_are_small & ( error_at_mean < zero) else: - use_symmetry_relation = (error_at_half < zero) + use_symmetry_relation = error_at_half < zero a_orig, y_orig = (a, y) a = tf.where(use_symmetry_relation, b, a) @@ -727,7 +727,7 @@ def root_finding_iteration(should_stop, low, high, candidate): halley_candidate)) new_delta = candidate - new_candidate - new_delta_is_negative = (new_delta < zero) + new_delta_is_negative = new_delta < zero new_low = tf.where(new_delta_is_negative, candidate, low) new_high = tf.where(new_delta_is_negative, high, candidate) @@ -782,7 +782,7 @@ def _betaincinv_partials(a, b, y, return_value=False): dtype_orig = dtype_util.common_dtype([a, b, y], tf.float32) # We promote bfloat16 and float16 to float32 to make this function consistent # with betaincinv. - should_promote_dtype = (dtype_orig in _f16bit_dtypes) + should_promote_dtype = dtype_orig in _f16bit_dtypes dtype = tf.float32 if should_promote_dtype else dtype_orig a, b, y = [tf.convert_to_tensor(z, dtype=dtype_orig) for z in [a, b, y]]