tf.keras.ops.numpy.logaddexp returns incorrect gradient at equal inputs

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf 2.21.0, also reproduced on tf 2.22.0-dev20260904

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

No response

Python version

Python 3.13.5

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

tf.keras.ops.numpy.logaddexp returns an incorrect first derivative at equal inputs.

The reproducer evaluates tf.keras.ops.numpy.logaddexp(t, [-4, 0, 5]) inside a weighted scalar expression at t = -4.0. In the first term, the two inputs to logaddexp are equal. Since log(exp(x) + exp(y)) is smooth at equality, the partial derivative with respect to the first input should be 1/2 at that point. A stabilization branch should preserve this derivative.

The forward value is computed correctly as -5.770614664699297, but reverse-mode autodiff returns 1.0359107226361899 for the derivative with respect to t. The expected derivative is approximately 0.53591072263619.

Expected behavior?

The first derivative of the weighted tf.keras.ops.numpy.logaddexp expression at t = -4.0 should be close to 0.53591072263619, not 1.0359107226361899.

Standalone code to reproduce the issue

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["TF_NUM_INTRAOP_THREADS"] = "1"
os.environ["TF_NUM_INTEROP_THREADS"] = "1"

import tensorflow as tf

def target(t):
    y = tf.keras.ops.numpy.logaddexp(
        t,
        tf.constant([-4, 0, 5], dtype=t.dtype),
    )
    return y[0] + 2 * y[1] - y[2] / 2

x = tf.constant(-4.0, dtype=tf.float64)

with tf.GradientTape() as tape:
    tape.watch(x)
    y = target(x)

actual = tape.gradient(y, x)
expected = 0.53591072263619

print("forward:", y.numpy())
print("actual:", actual.numpy())
print("expected:", expected)

if abs(float(actual) - expected) > 1e-9:
    print("BUG REPRODUCED: tf.keras.ops.numpy.logaddexp returns incorrect gradient at equal inputs")
else:
    print("not reproduced")

Relevant log output

forward: -5.770614664699297
actual: 1.0359107226361899
expected: 0.53591072263619
BUG REPRODUCED: tf.keras.ops.numpy.logaddexp returns incorrect gradient at equal inputs