Incorrect Gradient Computation when Nesting Zygote over TaylorDiff Description #92
Description
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