Skip to content

Commit 3d33ce2

Browse files
committed
Add new backends with DifferentiationInterface.jl
1 parent 1a5dd72 commit 3d33ce2

File tree

9 files changed

+331
-282
lines changed

9 files changed

+331
-282
lines changed

Project.toml

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,21 @@ version = "0.8.7"
44

55
[deps]
66
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
7+
DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
78
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
89
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
910
NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6"
10-
Requires = "ae029012-a4dd-5104-9daa-d747884805df"
1111
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
1212
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1313
SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
1414
SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
1515

1616
[compat]
17-
ADTypes = "1.2.1"
18-
ForwardDiff = "0.9.0, 0.10.0"
19-
NLPModels = "0.18, 0.19, 0.20, 0.21"
20-
Requires = "1"
21-
ReverseDiff = "1"
17+
ADTypes = "1.9.0"
18+
DifferentiationInterface = "0.6.0"
19+
ForwardDiff = "0.10.36"
20+
NLPModels = "0.21.3"
21+
ReverseDiff = "1.15.3"
2222
SparseConnectivityTracer = "0.6.1"
2323
SparseMatrixColorings = "0.4.0"
24-
julia = "^1.6"
24+
julia = "1.6"

docs/src/backend.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ The functions used internally to define the NLPModel API and the possible backen
1010
| Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
1111
| --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
1212
| `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeADGradient` | -- |
13-
| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | -- | `SparseADJacobian` |
14-
| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | -- | `SparseADHessian`/`SparseReverseADHessian` |
15-
| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | -- | -- |
16-
| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | -- | -- |
13+
| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `EnzymeADJacobian` | `SparseADJacobian` |
14+
| `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | -- | -- | `SparseADHessian`/`SparseReverseADHessian` |
15+
| `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeADJprod` | -- |
16+
| `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeADJtprod` | -- |
1717
| `Hvprod` | `ForwardDiffADHvprod`/`GenericForwardDiffADHvprod` | `ReverseDiffADHvprod`/`GenericReverseDiffADHvprod` | -- | -- | -- |
1818
| `directional_second_derivative` | `ForwardDiffADGHjvprod` | -- | -- | -- | -- |
1919

src/ADNLPModels.jl

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,17 @@ module ADNLPModels
44
using LinearAlgebra, SparseArrays
55

66
# external
7-
using ADTypes: ADTypes, AbstractColoringAlgorithm, AbstractSparsityDetector
7+
using ADTypes: ADTypes, AbstractADType, AbstractColoringAlgorithm, AbstractSparsityDetector, AutoEnzyme, AutoZygote
8+
using ADTypes: AutoForwardDiff, AutoReverseDiff, AutoMooncake, AutoDiffractor, AutoTracker, AutoSymbolics
9+
using ADTypes: AutoChainRules, AutoFastDifferentiation, AutoFiniteDiff, AutoFiniteDifferences, AutoPolyesterForwardDiff
10+
811
using SparseConnectivityTracer: TracerSparsityDetector
12+
import DifferentiationInterface
913
using SparseMatrixColorings
1014
using ForwardDiff, ReverseDiff
1115

1216
# JSO
1317
using NLPModels
14-
using Requires
1518

1619
abstract type AbstractADNLPModel{T, S} <: AbstractNLPModel{T, S} end
1720
abstract type AbstractADNLSModel{T, S} <: AbstractNLSModel{T, S} end
@@ -27,8 +30,7 @@ include("sparse_hessian.jl")
2730

2831
include("forward.jl")
2932
include("reverse.jl")
30-
include("enzyme.jl")
31-
include("zygote.jl")
33+
include("di.jl")
3234
include("predefined_backend.jl")
3335
include("nlp.jl")
3436

src/di.jl

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,254 @@
1+
for (ADGradient, fbackend) in ((:EnzymeADGradient , :AutoEnzyme ),
2+
(:ZygoteADGradient , :AutoZygote ),
3+
# (:ForwardDiffADGradient , :AutoForwardDiff ),
4+
# (:ReverseDiffADGradient , :AutoReverseDiff ),
5+
(:MooncakeADGradient , :AutoMooncake ),
6+
(:DiffractorADGradient , :AutoDiffractor ),
7+
(:TrackerADGradient , :AutoTracker ),
8+
(:SymbolicsADGradient , :AutoSymbolics ),
9+
(:ChainRulesADGradient , :AutoChainRules ),
10+
(:FastDifferentiationADGradient , :AutoFastDifferentiation ),
11+
(:FiniteDiffADGradient , :AutoFiniteDiff ),
12+
(:FiniteDifferencesADGradient , :AutoFiniteDifferences ),
13+
(:PolyesterForwardDiffADGradient, :AutoPolyesterForwardDiff))
14+
@eval begin
15+
16+
struct $ADGradient{B, E} <: ADBackend
17+
backend::B
18+
prep::E
19+
end
20+
21+
function $ADGradient(
22+
nvar::Integer,
23+
f,
24+
ncon::Integer = 0,
25+
c::Function = (args...) -> [];
26+
x0::AbstractVector = rand(nvar),
27+
kwargs...,
28+
)
29+
backend = $fbackend()
30+
prep = DifferentiationInterface.prepare_gradient(f, backend, x0)
31+
return $ADGradient(backend, prep)
32+
end
33+
34+
function gradient(b::$ADGradient, f, x)
35+
g = DifferentiationInterface.gradient(f, b.prep, b.backend, x)
36+
return g
37+
end
38+
39+
function gradient!(b::$ADGradient, g, f, x)
40+
DifferentiationInterface.gradient!(f, g, b.prep, b.backend, x)
41+
return g
42+
end
43+
44+
end
45+
end
46+
47+
for (ADJprod, fbackend) in ((:EnzymeADJprod , :AutoEnzyme ),
48+
(:ZygoteADJprod , :AutoZygote ),
49+
# (:ForwardDiffADJprod , :AutoForwardDiff ),
50+
# (:ReverseDiffADJprod , :AutoReverseDiff ),
51+
(:MooncakeADJprod , :AutoMooncake ),
52+
(:DiffractorADJprod , :AutoDiffractor ),
53+
(:TrackerADJprod , :AutoTracker ),
54+
(:SymbolicsADJprod , :AutoSymbolics ),
55+
(:ChainRulesADJprod , :AutoChainRules ),
56+
(:FastDifferentiationADJprod , :AutoFastDifferentiation ),
57+
(:FiniteDiffADJprod , :AutoFiniteDiff ),
58+
(:FiniteDifferencesADJprod , :AutoFiniteDifferences ),
59+
(:PolyesterForwardDiffADJprod, :AutoPolyesterForwardDiff))
60+
@eval begin
61+
62+
struct $ADJprod{B, E} <: ADBackend
63+
backend::B
64+
prep::E
65+
end
66+
67+
function $ADJprod(
68+
nvar::Integer,
69+
f,
70+
ncon::Integer = 0,
71+
c::Function = (args...) -> [];
72+
x0::AbstractVector = rand(nvar),
73+
kwargs...,
74+
)
75+
backend = $fbackend()
76+
dx = similar(x0, nvar)
77+
prep = DifferentiationInterface.prepare_pushforward(f, backend, x0, dx)
78+
return $ADJprod(backend, prep)
79+
end
80+
81+
function Jprod!(b::$ADJprod, Jv, f, x, v, ::Val)
82+
DifferentiationInterface.pushforward!(f, Jv, b.prep, b.backend, x, v)
83+
return Jv
84+
end
85+
86+
end
87+
end
88+
89+
for (ADJtprod, fbackend) in ((:EnzymeADJtprod , :AutoEnzyme ),
90+
(:ZygoteADJtprod , :AutoZygote ),
91+
# (:ForwardDiffADJtprod , :AutoForwardDiff ),
92+
# (:ReverseDiffADJtprod , :AutoReverseDiff ),
93+
(:MooncakeADJtprod , :AutoMooncake ),
94+
(:DiffractorADJtprod , :AutoDiffractor ),
95+
(:TrackerADJtprod , :AutoTracker ),
96+
(:SymbolicsADJtprod , :AutoSymbolics ),
97+
(:ChainRulesADJtprod , :AutoChainRules ),
98+
(:FastDifferentiationADJtprod , :AutoFastDifferentiation ),
99+
(:FiniteDiffADJtprod , :AutoFiniteDiff ),
100+
(:FiniteDifferencesADJtprod , :AutoFiniteDifferences ),
101+
(:PolyesterForwardDiffADJtprod, :AutoPolyesterForwardDiff))
102+
@eval begin
103+
104+
struct $ADJtprod{B, E} <: ADBackend
105+
backend::B
106+
prep::E
107+
end
108+
109+
function $ADJtprod(
110+
nvar::Integer,
111+
f,
112+
ncon::Integer = 0,
113+
c::Function = (args...) -> [];
114+
x0::AbstractVector = rand(nvar),
115+
kwargs...,
116+
)
117+
backend = $fbackend()
118+
dy = similar(x0, ncon)
119+
prep = DifferentiationInterface.prepare_pullback(f, backend, x0, dy)
120+
return $ADJtprod(backend, prep)
121+
end
122+
123+
function Jtprod!(b::$ADJtprod, Jtv, f, x, v, ::Val)
124+
DifferentiationInterface.pullback!(f, Jtv, b.prep, b.backend, x, v)
125+
return Jtv
126+
end
127+
128+
end
129+
end
130+
131+
for (ADJacobian, fbackend) in ((:EnzymeADJacobian , :AutoEnzyme ),
132+
(:ZygoteADJacobian , :AutoZygote ),
133+
# (:ForwardDiffADJacobian , :AutoForwardDiff ),
134+
# (:ReverseDiffADJacobian , :AutoReverseDiff ),
135+
(:MooncakeADJacobian , :AutoMooncake ),
136+
(:DiffractorADJacobian , :AutoDiffractor ),
137+
(:TrackerADJacobian , :AutoTracker ),
138+
(:SymbolicsADJacobian , :AutoSymbolics ),
139+
(:ChainRulesADJacobian , :AutoChainRules ),
140+
(:FastDifferentiationADJacobian , :AutoFastDifferentiation ),
141+
(:FiniteDiffADJacobian , :AutoFiniteDiff ),
142+
(:FiniteDifferencesADJacobian , :AutoFiniteDifferences ),
143+
(:PolyesterForwardDiffADJacobian, :AutoPolyesterForwardDiff))
144+
@eval begin
145+
146+
struct $ADJacobian{B, E} <: ADBackend
147+
backend::B
148+
prep::E
149+
end
150+
151+
function $ADJacobian(
152+
nvar::Integer,
153+
f,
154+
ncon::Integer = 0,
155+
c::Function = (args...) -> [];
156+
x0::AbstractVector = rand(nvar),
157+
kwargs...,
158+
)
159+
backend = $fbackend()
160+
y = similar(x0, ncon)
161+
prep = DifferentiationInterface.prepare_jacobian(f, y, backend, x0)
162+
return $ADJacobian(backend, prep)
163+
end
164+
165+
function jacobian(b::$ADJacobian, f, x)
166+
J = DifferentiationInterface.jacobian(f, b.prep, b.backend, x)
167+
return J
168+
end
169+
170+
end
171+
end
172+
173+
for (ADHvprod, fbackend) in ((:EnzymeADHvprod , :AutoEnzyme ),
174+
(:ZygoteADHvprod , :AutoZygote ),
175+
# (:ForwardDiffADHvprod , :AutoForwardDiff ),
176+
# (:ReverseDiffADHvprod , :AutoReverseDiff ),
177+
(:MooncakeADHvprod , :AutoMooncake ),
178+
(:DiffractorADHvprod , :AutoDiffractor ),
179+
(:TrackerADHvprod , :AutoTracker ),
180+
(:SymbolicsADHvprod , :AutoSymbolics ),
181+
(:ChainRulesADHvprod , :AutoChainRules ),
182+
(:FastDifferentiationADHvprod , :AutoFastDifferentiation ),
183+
(:FiniteDiffADHvprod , :AutoFiniteDiff ),
184+
(:FiniteDifferencesADHvprod , :AutoFiniteDifferences ),
185+
(:PolyesterForwardDiffADHvprod, :AutoPolyesterForwardDiff))
186+
@eval begin
187+
188+
struct $ADHvprod{B, E} <: ADBackend
189+
backend::B
190+
prep::E
191+
end
192+
193+
function $ADHvprod(
194+
nvar::Integer,
195+
f,
196+
ncon::Integer = 0,
197+
c::Function = (args...) -> [];
198+
x0::AbstractVector = rand(nvar),
199+
kwargs...,
200+
)
201+
backend = $fbackend()
202+
tx = similar(x0)
203+
prep = DifferentiationInterface.prepare_hvp(f, backend, x0, tx)
204+
return $ADHvprod(backend, prep)
205+
end
206+
207+
function Hvprod!(b::$ADHvprod, Hv, f, x, v, ::Val)
208+
DifferentiationInterface.hvp!(f, Hv, b.prep, b.backend, x, v)
209+
return Hv
210+
end
211+
212+
end
213+
end
214+
215+
for (ADHessian, fbackend) in ((:EnzymeADHessian , :AutoEnzyme ),
216+
(:ZygoteADHessian , :AutoZygote ),
217+
# (:ForwardDiffADHessian , :AutoForwardDiff ),
218+
# (:ReverseDiffADHessian , :AutoReverseDiff ),
219+
(:MooncakeADHessian , :AutoMooncake ),
220+
(:DiffractorADHessian , :AutoDiffractor ),
221+
(:TrackerADHessian , :AutoTracker ),
222+
(:SymbolicsADHessian , :AutoSymbolics ),
223+
(:ChainRulesADHessian , :AutoChainRules ),
224+
(:FastDifferentiationADHessian , :AutoFastDifferentiation ),
225+
(:FiniteDiffADHessian , :AutoFiniteDiff ),
226+
(:FiniteDifferencesADHessian , :AutoFiniteDifferences ),
227+
(:PolyesterForwardDiffADHessian, :AutoPolyesterForwardDiff))
228+
@eval begin
229+
230+
struct $ADHessian{B, E} <: ADBackend
231+
backend::B
232+
prep::E
233+
end
234+
235+
function $ADHessian(
236+
nvar::Integer,
237+
f,
238+
ncon::Integer = 0,
239+
c::Function = (args...) -> [];
240+
x0::AbstractVector = rand(nvar),
241+
kwargs...,
242+
)
243+
backend = $fbackend()
244+
prep = DifferentiationInterface.prepare_hessian(f, backend, x0)
245+
return $ADHessian(backend, prep)
246+
end
247+
248+
function hessian(b::$ADHessian, f, x)
249+
H = DifferentiationInterface.hessian(f, b.prep, b.backend, x)
250+
return H
251+
end
252+
253+
end
254+
end

src/enzyme.jl

Lines changed: 0 additions & 21 deletions
This file was deleted.

0 commit comments

Comments
 (0)