tf.math.cumulative_logsumexp produces NaN second derivative under nested forward-mode autodiff

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.math.cumulative_logsumexp produces nan for a second derivative computed with nested forward-mode autodiff.

The reproducer builds a finite float64 tensor from a scalar t, applies tf.math.cumulative_logsumexp with axis=-1, exclusive=False, and reverse=True, then reduces selected output entries into a scalar. With reverse=True, the selected outputs are smooth suffix log-sum-exp expressions, and their second derivatives are finite.

At t = -40.0, the forward value is computed correctly as 395.0000000000019, but nested tf.autodiff.ForwardAccumulator returns nan for the second derivative. The expected second derivative is approximately 1.0572349592992632e-12.

Expected behavior?

The second derivative of the weighted tf.math.cumulative_logsumexp output at t = -40.0 should be close to 1.0572349592992632e-12, not nan.

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):
    x = tf.reshape(
        tf.stack([t, -t + 1, t / 2 - 2, -2 * t, t + 3, t / 4]),
        [2, 3],
    )
    y = tf.math.cumulative_logsumexp(
        x,
        axis=-1,
        exclusive=False,
        reverse=True,
    )
    return y[0, 0] + 2 * y[0, 1] - y[0, 2] + 3 * y[1, 0] + y[1, 1] - 2 * y[1, 2]

def forward_ad(fn):
    def jvp(x):
        with tf.autodiff.ForwardAccumulator(x, tf.ones_like(x)) as acc:
            y = fn(x)
        return acc.jvp(y)
    return jvp

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

forward = target(x)
actual = forward_ad(forward_ad(target))(x)
expected = 1.0572349592992632e-12

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

if tf.math.is_nan(actual):
    print("BUG REPRODUCED: tf.math.cumulative_logsumexp produces NaN under nested forward AD")
else:
    print("not reproduced")

Relevant log output

forward: 395.0000000000019
actual: nan
expected: 1.0572349592992632e-12
BUG REPRODUCED: tf.math.cumulative_logsumexp produces NaN under nested forward AD