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