Incorrect Gradient Computation when Nesting Zygote over TaylorDiff Description

[mrazomej]

Bug Report: Incorrect Gradient Computation when Nesting Zygote over TaylorDiff

Description

When attempting to compute the gradient of a function that internally uses TaylorDiff for differentiation, Zygote returns incorrect results (all zeros).

Minimal Working Example

The following mwe demonstrates the issue:

import TaylorDiff
import Zygote

# Define a function
f(x) = sum(x .^ 2)

# Define function to return unit vectors
unit_vectors(x, i) = [j == i ? one(eltype(x)) : zero(eltype(x)) for j in 1:length(x)]

# Define input
x = Float64.(collect(1:5))

# Compute unit vectors for each direction to compute the gradient
e_vectors = [unit_vectors(x, i) for i in 1:length(x)]

# Define a function to compute gradient using TaylorDiff
∇f_taylor(f, x) = TaylorDiff.derivative.(Ref(f), Ref(x), e_vectors, Ref(Val(1)))

# Evaluate gradient
∇f_taylor(f, x) # Returns 5-element Vector{Float64}: 2.0 4.0 6.0 8.0 10.0 (correct)

# Define a simple function of the gradient with TaylorDiff
g_taylor(x) = sum(∇f_taylor(f, x))

# Evaluate gradient
g_taylor(x) # returns 30.0 (Correct)

# Compute gradient using Zygote
Zygote.gradient(g_taylor, x) # returns ([0.0, 0.0, 0.0, 0.0, 0.0],) (incorrect)

The function g_taylor(x) computes the sum of the gradient of f(x), which should return a nonzero result. Applying Zygote.gradient to g_taylor(x) should correctly compute the derivative.

Instead of returning the correct gradient, Zygote.gradient(g_taylor, x) returns all zeros.

Additional Information

Julia Version: 1.11.3 TaylorDiff Version: v0.3.1 Zygote Version: v0.6.75

[mrazomej]

Dear @tansongchen,

It's been a while since I posted this issue. The problem persists after the newest release of the package. Is there any chance this could be tackled relatively soon? Being able to do nested differentiation is one of the most attractive features of this package, and I am sure it would lead to a much higher adoption. Or, if it is not addressable, it would also be helpful to know.

On a similar note, is there a reason why the compatibility with Zygote is truncated at v0.6 when the most current version is v0.7.4?

Thanks in advance for your response.