JuliaDiff / JuliaDiff/TaylorDiff.jl

Incorrect Gradient Computation when Nesting Zygote over TaylorDiff Description

Open
#92 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
83
Forks
12
PR merge metrics
No merged PRs in 30d

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

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the minimal example with the reported Julia, TaylorDiff, and Zygote versions, focusing on TaylorDiff.derivative nested inside Zygote.gradient. Trace the interaction between those entry points and add a regression test for the example. Done means the computed gradient is no longer all zeros and matches the expected nonzero result.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
devtools
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.