Fix NaN gradient of tf.math.igamma at a=1, x=0 - #123580

Code Review

This pull request replaces (a - 1) * log(x) with xlogy(a - 1, x) in the gradient computation of igamma (_IgammaGrad) to prevent NaN values when a == 1 and x == 0, and adds corresponding unit tests. The reviewer feedback suggests applying the same xlogy fix to the complementary incomplete gamma function gradient (_IgammacGrad) and updating the new unit tests to run in both graph and eager modes using tf.GradientTape instead of being restricted to V1 graph mode.

Thank you for this contribution. The fix for the NaN gradient in tf.math.igamma at $a=1, x=0$ looks mathematically sound and well-targeted.

Analysis

The original formulation in _IgammaGrad computed partial_x = exp(-x + (a - 1) * log(x) - lgamma(a)).
At $a=1, x=0$, the term (a - 1) * log(x) evaluates to 0 * -inf = NaN, despite the true mathematical derivative at that point being $e^{-0} = 1$.

Replacing (a - 1) * log(x) with xlogy(a - 1, x) robustly handles this edge case because xlogy natively maps $0 \times \log(0)$ to $0$.

Regimes Verification

• $a > 1, x = 0$: xlogy(a-1, 0) evaluates to $-\infty$ (since $a-1 > 0$). The exponent becomes $-\infty$, and $e^{-\infty} = 0$. This is correct as $x^{a-1} = 0$ for $a > 1$.
• $a < 1, x = 0$: xlogy(a-1, 0) evaluates to $+\infty$ (since $a-1 < 0$). The exponent becomes $+\infty$, and $e^{+\infty} = +\infty$. This is correct as $x^{a-1} = \infty$ for $a < 1$.
• $x > 0$: xlogy(a-1, x) is identical to $(a-1)\log(x)$.

Unit Tests

We verified this behavior with the following unit test:

import tensorflow as tf
from absl.testing import parameterized

class IgammaGradTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.parameters(tf.float32, tf.float64)
  def testNaNEdgeCase(self, dtype):
    a = tf.constant(1.0, dtype=dtype)
    x = tf.constant(0.0, dtype=dtype)
    with tf.GradientTape() as tape:
      tape.watch(x)
      y = tf.math.igamma(a, x)
    grad = tape.gradient(y, x)
    self.assertAllClose(grad, 1.0)

This ensures both single and double precision backends behave consistently.