Zygote: Differentiate over Zygote operation?

Specific Domains Numerics
, ,
1

Hello everyone,

I wondered, if it is possible to differentiate with Zygote over another Zygote-Operation (I attached a minimal example).

using Zygote

function myFunction(a)
    Zygote.jacobian(sin, a)[1]
end

Zygote.jacobian(myFunction, 0.5)

At this point, I know that there are other ways to build hessians, this is just a simple code example. In fact I want to build derivatives that are more complicated. Code example generates the error message:

ERROR: LoadError: UndefVarError: xs not defined
Stacktrace:
[1] (::Zygote.var"#442#443")(#unused#::Nothing)
@ Zygote ….julia\packages\Zygote\TaBlo\src\lib\array.jl:79
[2] (::Zygote.var"#2359#back#444"{Zygote.var"#442#443"})(Δ::Nothing)
@ Zygote ….julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67
[3] Pullback
@ ….julia\packages\Zygote\TaBlo\src\lib\grad.jl:184 [inlined]
[4] (::typeof(∂(_gradcopy!)))(Δ::Nothing)
@ Zygote ….julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
[5] Pullback
@ ….julia\packages\Zygote\TaBlo\src\lib\grad.jl:165 [inlined]
[6] (::typeof(∂(withjacobian)))(Δ::NamedTuple{(:val, :grad), Tuple{Nothing, Tuple{Vector{Float64}}}})
@ Zygote ….julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
[7] (::Zygote.var"#209#210"{Tuple{Tuple{Nothing}, Tuple{Nothing}}, typeof(∂(withjacobian))})(Δ::NamedTuple{(:val, :grad), Tuple{Nothing, Tuple{Vector{Float64}}}})
@ Zygote ….julia\packages\Zygote\TaBlo\src\lib\lib.jl:203
[8] #1754#back
@ ….julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67 [inlined]
[9] Pullback
@ ….julia\packages\Zygote\TaBlo\src\lib\grad.jl:140 [inlined]
[10] (::typeof(∂(jacobian)))(Δ::Tuple{Vector{Float64}})
@ Zygote ….julia\packages\Zygote\TaBlo\src\compiler\interface2.jl:0
[11] Pullback
@ …\Julia\zygoteRecursion.jl:7 [inlined]
[12] #209
@ ….julia\packages\Zygote\TaBlo\src\lib\lib.jl:203 [inlined]
[13] #1754#back
@ ….julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67 [inlined]
[14] Pullback
@ .\operators.jl:938 [inlined]
[15] (::Zygote.var"#46#47"{typeof(∂(ComposedFunction{typeof(Zygote._jvec), typeof(myFunction)}(Zygote._jvec, myFunction)))})(Δ::Vector{Float64})
@ Zygote ….julia\packages\Zygote\TaBlo\src\compiler\interface.jl:41
[16] withjacobian(f::Function, args::Float64)
@ Zygote ….julia\packages\Zygote\TaBlo\src\lib\grad.jl:162
[17] jacobian(f::Function, args::Float64)
@ Zygote ….julia\packages\Zygote\TaBlo\src\lib\grad.jl:140
[18] top-level scope
@ …\Julia\zygoteRecursion.jl:13
in expression starting at …\Julia\zygoteRecursion.jl:13

[Julia v1.6, Zygote v0.6.17]

So is there a way to differentiate over other Zygote-differentiation-calls?

Thanks in advance!

2

I get a different error trace than you get. However, if I redefine myFunction(a) as

function myFunction(a)
    Zygote.gradient(sin, a)[1]
end

The whole thing works. Maybe this is a bug / missing definition of Zygote?

1 Like
3

Yes, “bug / missing definition” is about right.

Zygote’s jacobian works by making the output array dx (one for each argument) and then writing slices into it, here, which won’t work within a derivative. It would not be very hard to avoid that, either by gathering all the slices and combining them using a differentiable function, or by adding another rule for the this part of the operation. Someone would just have to get around to doing it.

This is the story in general for 2nd derivatives. The machinery does in principle work, but particular rules or functions don’t co-operate, and would need to be fixed.

There are very few tests involving second derivatives. An easier way to contribute than upgrading things would be to invent many test cases of things which do now work, and use various different functions, and add those to the package’s tests so that they are not accidentally broken in the process of fixing other things.

1 Like
4

Didn’t try that one at first, but you are right, Gradients are working!

5

Thank you for the reply. I just tried to get this working with the minimal example from above (in this case there is only one argument, so the jacobian is a tuple containing a single vector containing a single value)… but without success. Could you provide a working example? That would be helpful :slight_smile: Thanks in advance!