EnzymeAD / EnzymeAD/Enzyme.jl

Example of JVP / J'VP with Krylov.jl

Open
#1,957 6 comments 0 reactions 0 assignees View on GitHub
documentation
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

cc: @michel2323 & @amontoison

With @lcandiot I was wondering how to write a proper JVP J'VP with Enzyme and finally converged.

This might turn into a nice example one of these days.
@wsmoses any ideas on how to avoid the calls to `zero(y)` and`zero(w)`/`copy(w)`?

```julia
using Krylov, Enzyme, LinearOperators, ForwardDiff, LinearAlgebra

xk = ones(2)

F(x) = [x[1]^4 - 3; exp(x[2]) - 2; log(x[1]) - x[2]^2]

function JVP!(y, f::F, u, v) where F
Enzyme.autodiff(Forward,
(temp, v) -> (temp .= f(v); nothing),
Const,
DuplicatedNoNeed(zero(y), y),
DuplicatedNoNeed(u, v))
return nothing
end

"""
Calculate the Jacobian-Transpose Vector Product in-place by updating `y`.
"""
function JᵀVP!(y, f::F, u, w) where F
y .= 0 # Enzyme expects y to be zero
Enzyme.autodiff(Enzyme.Reverse,
(out, x) -> (out .= f(x); nothing),
Const,
DuplicatedNoNeed(zero(w), copy(w)), # copy since otherwise Enzyme will zero
DuplicatedNoNeed(u, y))
return nothing
end

J(y, v) = ForwardDiff.derivative!(y, t -> F(xk + t * v), 0)
Jᵀ(y, u) = ForwardDiff.gradient!(y, x -> dot(F(x), u), xk)

w = rand(3)
v = rand(2)

y_fwd = zeros(2)
Jᵀ(y_fwd, w)
@show y_fwd

y_enz = zeros(2)
@show JᵀVP!(y_enz, F, xk, w)
@show y_enz

@assert y_enz ≈ y_fwd

y2_fwd = zeros(3)
J(y2_fwd, v)
@show y2_fwd

y2_enz = zeros(3)
@show JVP!(y2_enz, F, xk, v)
@show y2_enz

@assert y2_enz ≈ y2_fwd

opJ_FWD = LinearOperator(Float64, 3, 2, false, false, (y, v) -> J(y, v),
(y, w) -> Jᵀ(y, w),
(y, u) -> Jᵀ(y, u))

x_forward, _ = lsmr(opJ_FWD, -F(xk))

opJ = LinearOperator(Float64, 3, 2, false, false, (y, v) -> JVP!(y, F, xk, v),
(y, w) -> JᵀVP!(y,F, xk, w),
(y, u) -> JᵀVP!(y,F, xk, u))

x_enzyme, _ = lsmr(opJ, -F(xk))
x_enzyme ≈ x_forward
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.