Skip to content

Commit 13d863b

Browse files
prbzrgprbzrg
committed
refactor jacobian parts
1 parent 51e57fa commit 13d863b

2 files changed

Lines changed: 132 additions & 83 deletions

File tree

‎src/icnf.jl‎

Lines changed: 16 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ function augmented_f(
120120
n_aug = n_augment(icnf, mode)
121121
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
122122
z = u[begin:(end - n_aug - 1)]
123-
ż, J = DifferentiationInterface.value_and_jacobian(snn, icnf.compute_mode.adback, z)
123+
ż, J = icnf_jacobian(icnf, mode, snn, z)
124124
l̇ = -LinearAlgebra.tr(J)
125125
return vcat(ż, l̇)
126126
end
@@ -139,7 +139,7 @@ function augmented_f(
139139
n_aug = n_augment(icnf, mode)
140140
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
141141
z = u[begin:(end - n_aug - 1)]
142-
ż, J = DifferentiationInterface.value_and_jacobian(snn, icnf.compute_mode.adback, z)
142+
ż, J = icnf_jacobian(icnf, mode, snn, z)
143143
du[begin:(end - n_aug - 1)] .= ż
144144
du[(end - n_aug)] = -LinearAlgebra.tr(J)
145145
return nothing
@@ -158,7 +158,7 @@ function augmented_f(
158158
n_aug = n_augment(icnf, mode)
159159
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
160160
z = u[begin:(end - n_aug - 1), :]
161-
ż, J = jacobian_batched(icnf, snn, z)
161+
ż, J = icnf_jacobian(icnf, mode, snn, z)
162162
l̇ = -transpose(LinearAlgebra.tr.(J))
163163
return vcat(ż, l̇)
164164
end
@@ -177,7 +177,7 @@ function augmented_f(
177177
n_aug = n_augment(icnf, mode)
178178
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
179179
z = u[begin:(end - n_aug - 1), :]
180-
ż, J = jacobian_batched(icnf, snn, z)
180+
ż, J = icnf_jacobian(icnf, mode, snn, z)
181181
du[begin:(end - n_aug - 1), :] .= ż
182182
du[(end - n_aug), :] .= -(LinearAlgebra.tr.(J))
183183
return nothing
@@ -196,9 +196,7 @@ function augmented_f(
196196
n_aug = n_augment(icnf, mode)
197197
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
198198
z = u[begin:(end - n_aug - 1)]
199-
ż, ϵJ =
200-
DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
201-
ϵJ = only(ϵJ)
199+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
202200
l̇ = -LinearAlgebra.dot(ϵJ, ϵ)
203201
Ė = if NORM_Z
204202
LinearAlgebra.norm(ż)
@@ -227,9 +225,7 @@ function augmented_f(
227225
n_aug = n_augment(icnf, mode)
228226
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
229227
z = u[begin:(end - n_aug - 1)]
230-
ż, ϵJ =
231-
DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
232-
ϵJ = only(ϵJ)
228+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
233229
du[begin:(end - n_aug - 1)] .= ż
234230
du[(end - n_aug)] = -LinearAlgebra.dot(ϵJ, ϵ)
235231
du[(end - n_aug + 1)] = if NORM_Z
@@ -258,13 +254,7 @@ function augmented_f(
258254
n_aug = n_augment(icnf, mode)
259255
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
260256
z = u[begin:(end - n_aug - 1)]
261-
ż, Jϵ = DifferentiationInterface.value_and_pushforward(
262-
snn,
263-
icnf.compute_mode.adback,
264-
z,
265-
(ϵ,),
266-
)
267-
Jϵ = only(Jϵ)
257+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
268258
l̇ = -LinearAlgebra.dot(ϵ, Jϵ)
269259
Ė = if NORM_Z
270260
LinearAlgebra.norm(ż)
@@ -293,13 +283,7 @@ function augmented_f(
293283
n_aug = n_augment(icnf, mode)
294284
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
295285
z = u[begin:(end - n_aug - 1)]
296-
ż, Jϵ = DifferentiationInterface.value_and_pushforward(
297-
snn,
298-
icnf.compute_mode.adback,
299-
z,
300-
(ϵ,),
301-
)
302-
Jϵ = only(Jϵ)
286+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
303287
du[begin:(end - n_aug - 1)] .= ż
304288
du[(end - n_aug)] = -LinearAlgebra.dot(ϵ, Jϵ)
305289
du[(end - n_aug + 1)] = if NORM_Z
@@ -328,9 +312,7 @@ function augmented_f(
328312
n_aug = n_augment(icnf, mode)
329313
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
330314
z = u[begin:(end - n_aug - 1), :]
331-
ż, ϵJ =
332-
DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
333-
ϵJ = only(ϵJ)
315+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
334316
l̇ = -sum(ϵJ .* ϵ; dims = 1)
335317
Ė = transpose(if NORM_Z
336318
LinearAlgebra.norm.(eachcol(ż))
@@ -363,9 +345,7 @@ function augmented_f(
363345
n_aug = n_augment(icnf, mode)
364346
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
365347
z = u[begin:(end - n_aug - 1), :]
366-
ż, ϵJ =
367-
DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
368-
ϵJ = only(ϵJ)
348+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
369349
du[begin:(end - n_aug - 1), :] .= ż
370350
du[(end - n_aug), :] .= -vec(sum(ϵJ .* ϵ; dims = 1))
371351
du[(end - n_aug + 1), :] .= if NORM_Z
@@ -394,13 +374,7 @@ function augmented_f(
394374
n_aug = n_augment(icnf, mode)
395375
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
396376
z = u[begin:(end - n_aug - 1), :]
397-
ż, Jϵ = DifferentiationInterface.value_and_pushforward(
398-
snn,
399-
icnf.compute_mode.adback,
400-
z,
401-
(ϵ,),
402-
)
403-
Jϵ = only(Jϵ)
377+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
404378
l̇ = -sum(ϵ .* Jϵ; dims = 1)
405379
Ė = transpose(if NORM_Z
406380
LinearAlgebra.norm.(eachcol(ż))
@@ -433,13 +407,7 @@ function augmented_f(
433407
n_aug = n_augment(icnf, mode)
434408
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
435409
z = u[begin:(end - n_aug - 1), :]
436-
ż, Jϵ = DifferentiationInterface.value_and_pushforward(
437-
snn,
438-
icnf.compute_mode.adback,
439-
z,
440-
(ϵ,),
441-
)
442-
Jϵ = only(Jϵ)
410+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
443411
du[begin:(end - n_aug - 1), :] .= ż
444412
du[(end - n_aug), :] .= -vec(sum(ϵ .* Jϵ; dims = 1))
445413
du[(end - n_aug + 1), :] .= if NORM_Z
@@ -468,8 +436,7 @@ function augmented_f(
468436
n_aug = n_augment(icnf, mode)
469437
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
470438
z = u[begin:(end - n_aug - 1), :]
471-
ż = snn(z)
472-
ϵJ = Lux.vector_jacobian_product(snn, icnf.compute_mode.adback, z, ϵ)
439+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
473440
l̇ = -sum(ϵJ .* ϵ; dims = 1)
474441
Ė = transpose(if NORM_Z
475442
LinearAlgebra.norm.(eachcol(ż))
@@ -502,8 +469,7 @@ function augmented_f(
502469
n_aug = n_augment(icnf, mode)
503470
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
504471
z = u[begin:(end - n_aug - 1), :]
505-
ż = snn(z)
506-
ϵJ = Lux.vector_jacobian_product(snn, icnf.compute_mode.adback, z, ϵ)
472+
ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
507473
du[begin:(end - n_aug - 1), :] .= ż
508474
du[(end - n_aug), :] .= -vec(sum(ϵJ .* ϵ; dims = 1))
509475
du[(end - n_aug + 1), :] .= if NORM_Z
@@ -532,8 +498,7 @@ function augmented_f(
532498
n_aug = n_augment(icnf, mode)
533499
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
534500
z = u[begin:(end - n_aug - 1), :]
535-
ż = snn(z)
536-
Jϵ = Lux.jacobian_vector_product(snn, icnf.compute_mode.adback, z, ϵ)
501+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
537502
l̇ = -sum(ϵ .* Jϵ; dims = 1)
538503
Ė = transpose(if NORM_Z
539504
LinearAlgebra.norm.(eachcol(ż))
@@ -566,8 +531,7 @@ function augmented_f(
566531
n_aug = n_augment(icnf, mode)
567532
snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
568533
z = u[begin:(end - n_aug - 1), :]
569-
ż = snn(z)
570-
Jϵ = Lux.jacobian_vector_product(snn, icnf.compute_mode.adback, z, ϵ)
534+
ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
571535
du[begin:(end - n_aug - 1), :] .= ż
572536
du[(end - n_aug), :] .= -vec(sum(ϵ .* Jϵ; dims = 1))
573537
du[(end - n_aug + 1), :] .= if NORM_Z

‎src/utils.jl‎

Lines changed: 116 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,57 +1,152 @@
1-
function jacobian_batched(
1+
function icnf_jacobian(
2+
icnf::AbstractICNF{<:AbstractFloat, <:DIVectorMode},
3+
::TestMode,
4+
f::LuxCore.StatefulLuxLayer,
5+
xs::AbstractVector{<:Real},
6+
)
7+
y, J = DifferentiationInterface.value_and_jacobian(f, icnf.compute_mode.adback, xs)
8+
return y, oftype(y, J)
9+
end
10+
11+
function icnf_jacobian(
12+
icnf::AbstractICNF{<:AbstractFloat, <:DIMatrixMode},
13+
::TestMode,
14+
f::LuxCore.StatefulLuxLayer,
15+
xs::AbstractMatrix{<:Real},
16+
)
17+
y, J = DifferentiationInterface.value_and_jacobian(f, icnf.compute_mode.adback, xs)
18+
return y, oftype.(Ref(y), split_jac(J, size(xs, 1)))
19+
end
20+
21+
function icnf_jacobian(
222
icnf::AbstractICNF{T, <:DIVecJacMatrixMode},
23+
::TestMode,
324
f::LuxCore.StatefulLuxLayer,
425
xs::AbstractMatrix{<:Real},
526
) where {T}
627
y = f(xs)
728
z = similar(xs)
829
ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
9-
res = Zygote.Buffer(
10-
convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
11-
size(xs, 1),
12-
size(xs, 1),
13-
size(xs, 2),
14-
)
30+
res = Zygote.Buffer(y, size(xs, 1), size(xs, 1), size(xs, 2))
1531
for i in axes(xs, 1)
1632
ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
1733
res[i, :, :] =
1834
only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (z,)))
1935
ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
2036
end
21-
return y, eachslice(copy(res); dims = 3)
37+
return y, oftype.(Ref(y), eachslice(copy(res); dims = 3))
2238
end
2339

24-
function jacobian_batched(
40+
function icnf_jacobian(
2541
icnf::AbstractICNF{T, <:DIJacVecMatrixMode},
42+
::TestMode,
2643
f::LuxCore.StatefulLuxLayer,
2744
xs::AbstractMatrix{<:Real},
2845
) where {T}
2946
y = f(xs)
3047
z = similar(xs)
3148
ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
32-
res = Zygote.Buffer(
33-
convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
34-
size(xs, 1),
35-
size(xs, 1),
36-
size(xs, 2),
37-
)
49+
res = Zygote.Buffer(y, size(xs, 1), size(xs, 1), size(xs, 2))
3850
for i in axes(xs, 1)
3951
ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
4052
res[:, i, :] = only(
4153
DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (z,)),
4254
)
4355
ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
4456
end
45-
return y, eachslice(copy(res); dims = 3)
57+
return y, oftype.(Ref(y), eachslice(copy(res); dims = 3))
4658
end
4759

48-
function jacobian_batched(
49-
icnf::AbstractICNF{T, <:DIMatrixMode},
60+
function icnf_jacobian(
61+
icnf::AbstractICNF{<:AbstractFloat, <:LuxMatrixMode},
62+
::TestMode,
5063
f::LuxCore.StatefulLuxLayer,
5164
xs::AbstractMatrix{<:Real},
52-
) where {T}
53-
y, J = DifferentiationInterface.value_and_jacobian(f, icnf.compute_mode.adback, xs)
54-
return y, split_jac(J, size(xs, 1))
65+
)
66+
y = f(xs)
67+
return y,
68+
oftype.(
69+
Ref(y),
70+
eachslice(Lux.batched_jacobian(f, icnf.compute_mode.adback, xs); dims = 3),
71+
)
72+
end
73+
74+
function icnf_jacobian(
75+
icnf::AbstractICNF{T, <:DIVecJacVectorMode},
76+
::TrainMode,
77+
f::LuxCore.StatefulLuxLayer,
78+
xs::AbstractVector{<:Real},
79+
ϵ::AbstractVector{T},
80+
) where {T <: AbstractFloat}
81+
y, ϵJ =
82+
DifferentiationInterface.value_and_pullback(f, icnf.compute_mode.adback, xs, (ϵ,))
83+
return y, oftype(y, only(ϵJ))
84+
end
85+
86+
function icnf_jacobian(
87+
icnf::AbstractICNF{T, <:DIJacVecVectorMode},
88+
::TrainMode,
89+
f::LuxCore.StatefulLuxLayer,
90+
xs::AbstractVector{<:Real},
91+
ϵ::AbstractVector{T},
92+
) where {T <: AbstractFloat}
93+
y, Jϵ = DifferentiationInterface.value_and_pushforward(
94+
f,
95+
icnf.compute_mode.adback,
96+
xs,
97+
(ϵ,),
98+
)
99+
return y, oftype(y, only(Jϵ))
100+
end
101+
102+
function icnf_jacobian(
103+
icnf::AbstractICNF{T, <:DIVecJacMatrixMode},
104+
::TrainMode,
105+
f::LuxCore.StatefulLuxLayer,
106+
xs::AbstractMatrix{<:Real},
107+
ϵ::AbstractMatrix{T},
108+
) where {T <: AbstractFloat}
109+
y, ϵJ =
110+
DifferentiationInterface.value_and_pullback(f, icnf.compute_mode.adback, xs, (ϵ,))
111+
return y, oftype(y, only(ϵJ))
112+
end
113+
114+
function icnf_jacobian(
115+
icnf::AbstractICNF{T, <:DIJacVecMatrixMode},
116+
::TrainMode,
117+
f::LuxCore.StatefulLuxLayer,
118+
xs::AbstractMatrix{<:Real},
119+
ϵ::AbstractMatrix{T},
120+
) where {T <: AbstractFloat}
121+
y, Jϵ = DifferentiationInterface.value_and_pushforward(
122+
f,
123+
icnf.compute_mode.adback,
124+
xs,
125+
(ϵ,),
126+
)
127+
return y, oftype(y, only(Jϵ))
128+
end
129+
130+
function icnf_jacobian(
131+
icnf::AbstractICNF{T, <:LuxVecJacMatrixMode},
132+
::TrainMode,
133+
f::LuxCore.StatefulLuxLayer,
134+
xs::AbstractMatrix{<:Real},
135+
ϵ::AbstractMatrix{T},
136+
) where {T <: AbstractFloat}
137+
y = f(xs)
138+
return y, oftype(y, Lux.vector_jacobian_product(f, icnf.compute_mode.adback, xs, ϵ))
139+
end
140+
141+
function icnf_jacobian(
142+
icnf::AbstractICNF{T, <:LuxJacVecMatrixMode},
143+
::TrainMode,
144+
f::LuxCore.StatefulLuxLayer,
145+
xs::AbstractMatrix{<:Real},
146+
ϵ::AbstractMatrix{T},
147+
) where {T <: AbstractFloat}
148+
y = f(xs)
149+
return y, oftype(y, Lux.jacobian_vector_product(f, icnf.compute_mode.adback, xs, ϵ))
55150
end
56151

57152
function split_jac(x::AbstractMatrix{<:Real}, sz::Integer)
@@ -62,13 +157,3 @@ function split_jac(x::AbstractMatrix{<:Real}, sz::Integer)
62157
)
63158
)
64159
end
65-
66-
function jacobian_batched(
67-
icnf::AbstractICNF{T, <:LuxMatrixMode},
68-
f::LuxCore.StatefulLuxLayer,
69-
xs::AbstractMatrix{<:Real},
70-
) where {T}
71-
y = f(xs)
72-
J = Lux.batched_jacobian(f, icnf.compute_mode.adback, xs)
73-
return y, eachslice(J; dims = 3)
74-
end

0 commit comments

Comments
 (0)