refactor jacobian parts

.github/workflows/CI-CheckBy.yml

@@ -26,7 +26,7 @@ jobs:
         version:
           - release
           - lts
-          - nightly
+          # - nightly
         os:
           - ubuntu-latest
           # - macOS-latest

src/icnf.jl

@@ -120,7 +120,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, J = DifferentiationInterface.value_and_jacobian(snn, icnf.compute_mode.adback, z)
+    ż, J = icnf_jacobian(icnf, mode, snn, z)
     l̇ = -LinearAlgebra.tr(J)
     return vcat(ż, l̇)
 end
@@ -139,7 +139,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, J = DifferentiationInterface.value_and_jacobian(snn, icnf.compute_mode.adback, z)
+    ż, J = icnf_jacobian(icnf, mode, snn, z)
     du[begin:(end - n_aug - 1)] .= ż
     du[(end - n_aug)] = -LinearAlgebra.tr(J)
     return nothing
@@ -158,8 +158,8 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, J = jacobian_batched(icnf, snn, z)
-    l̇ = -transpose(LinearAlgebra.tr.(J))
+    ż, J = icnf_jacobian(icnf, mode, snn, z)
+    l̇ = -transpose(LinearAlgebra.tr.(eachslice(J; dims = 3)))
     return vcat(ż, l̇)
 end
 
@@ -177,9 +177,9 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, J = jacobian_batched(icnf, snn, z)
+    ż, J = icnf_jacobian(icnf, mode, snn, z)
     du[begin:(end - n_aug - 1), :] .= ż
-    du[(end - n_aug), :] .= -(LinearAlgebra.tr.(J))
+    du[(end - n_aug), :] .= -(LinearAlgebra.tr.(eachslice(J; dims = 3)))
     return nothing
 end
 
@@ -196,9 +196,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, ϵJ =
-        DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
-    ϵJ = only(ϵJ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -LinearAlgebra.dot(ϵJ, ϵ)
     Ė = if NORM_Z
         LinearAlgebra.norm(ż)
@@ -227,9 +225,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, ϵJ =
-        DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
-    ϵJ = only(ϵJ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1)] .= ż
     du[(end - n_aug)] = -LinearAlgebra.dot(ϵJ, ϵ)
     du[(end - n_aug + 1)] = if NORM_Z
@@ -258,13 +254,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, Jϵ = DifferentiationInterface.value_and_pushforward(
-        snn,
-        icnf.compute_mode.adback,
-        z,
-        (ϵ,),
-    )
-    Jϵ = only(Jϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -LinearAlgebra.dot(ϵ, Jϵ)
     Ė = if NORM_Z
         LinearAlgebra.norm(ż)
@@ -293,13 +283,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1)]
-    ż, Jϵ = DifferentiationInterface.value_and_pushforward(
-        snn,
-        icnf.compute_mode.adback,
-        z,
-        (ϵ,),
-    )
-    Jϵ = only(Jϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1)] .= ż
     du[(end - n_aug)] = -LinearAlgebra.dot(ϵ, Jϵ)
     du[(end - n_aug + 1)] = if NORM_Z
@@ -328,9 +312,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, ϵJ =
-        DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
-    ϵJ = only(ϵJ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -sum(ϵJ .* ϵ; dims = 1)
     Ė = transpose(if NORM_Z
         LinearAlgebra.norm.(eachcol(ż))
@@ -363,9 +345,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, ϵJ =
-        DifferentiationInterface.value_and_pullback(snn, icnf.compute_mode.adback, z, (ϵ,))
-    ϵJ = only(ϵJ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1), :] .= ż
     du[(end - n_aug), :] .= -vec(sum(ϵJ .* ϵ; dims = 1))
     du[(end - n_aug + 1), :] .= if NORM_Z
@@ -394,13 +374,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, Jϵ = DifferentiationInterface.value_and_pushforward(
-        snn,
-        icnf.compute_mode.adback,
-        z,
-        (ϵ,),
-    )
-    Jϵ = only(Jϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -sum(ϵ .* Jϵ; dims = 1)
     Ė = transpose(if NORM_Z
         LinearAlgebra.norm.(eachcol(ż))
@@ -433,13 +407,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż, Jϵ = DifferentiationInterface.value_and_pushforward(
-        snn,
-        icnf.compute_mode.adback,
-        z,
-        (ϵ,),
-    )
-    Jϵ = only(Jϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1), :] .= ż
     du[(end - n_aug), :] .= -vec(sum(ϵ .* Jϵ; dims = 1))
     du[(end - n_aug + 1), :] .= if NORM_Z
@@ -468,8 +436,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż = snn(z)
-    ϵJ = Lux.vector_jacobian_product(snn, icnf.compute_mode.adback, z, ϵ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -sum(ϵJ .* ϵ; dims = 1)
     Ė = transpose(if NORM_Z
         LinearAlgebra.norm.(eachcol(ż))
@@ -502,8 +469,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż = snn(z)
-    ϵJ = Lux.vector_jacobian_product(snn, icnf.compute_mode.adback, z, ϵ)
+    ż, ϵJ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1), :] .= ż
     du[(end - n_aug), :] .= -vec(sum(ϵJ .* ϵ; dims = 1))
     du[(end - n_aug + 1), :] .= if NORM_Z
@@ -532,8 +498,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż = snn(z)
-    Jϵ = Lux.jacobian_vector_product(snn, icnf.compute_mode.adback, z, ϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     l̇ = -sum(ϵ .* Jϵ; dims = 1)
     Ė = transpose(if NORM_Z
         LinearAlgebra.norm.(eachcol(ż))
@@ -566,8 +531,7 @@ function augmented_f(
     n_aug = n_augment(icnf, mode)
     snn = LuxCore.StatefulLuxLayer{true}(nn, p, st)
     z = u[begin:(end - n_aug - 1), :]
-    ż = snn(z)
-    Jϵ = Lux.jacobian_vector_product(snn, icnf.compute_mode.adback, z, ϵ)
+    ż, Jϵ = icnf_jacobian(icnf, mode, snn, z, ϵ)
     du[begin:(end - n_aug - 1), :] .= ż
     du[(end - n_aug), :] .= -vec(sum(ϵ .* Jϵ; dims = 1))
     du[(end - n_aug + 1), :] .= if NORM_Z

src/utils.jl

@@ -1,74 +1,165 @@
-function jacobian_batched(
+function icnf_jacobian(
+    icnf::AbstractICNF{<:AbstractFloat, <:DIVectorMode},
+    ::TestMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractVector{<:Real},
+)
+    y = f(xs)
+    return y,
+    oftype(hcat(y), DifferentiationInterface.jacobian(f, icnf.compute_mode.adback, xs))
+end
+
+function icnf_jacobian(
+    icnf::AbstractICNF{<:AbstractFloat, <:DIMatrixMode},
+    ::TestMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractMatrix{<:Real},
+)
+    y = f(xs)
+    J = DifferentiationInterface.jacobian(f, icnf.compute_mode.adback, xs)
+    return y,
+    oftype(
+        cat(y; dims = Val(3)),
+        cat(
+            (
+                J[i:j, i:j] for (i, j) in zip(
+                    firstindex(J, 1):size(y, 1):lastindex(J, 1),
+                    (firstindex(J, 1) + size(y, 1) - 1):size(y, 1):lastindex(J, 1),
+                )
+            )...;
+            dims = Val(3),
+        ),
+    )
+end
+
+function icnf_jacobian(
     icnf::AbstractICNF{T, <:DIVecJacMatrixMode},
+    ::TestMode,
     f::LuxCore.StatefulLuxLayer,
     xs::AbstractMatrix{<:Real},
-) where {T}
+) where {T <: AbstractFloat}
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(
-        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
-        size(xs, 1),
-        size(xs, 1),
-        size(xs, 2),
-    )
+    res = Zygote.Buffer(y, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[i, :, :] =
             only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (z,)))
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
     end
-    return y, eachslice(copy(res); dims = 3)
+    return y, oftype(cat(y; dims = Val(3)), copy(res))
 end
 
-function jacobian_batched(
+function icnf_jacobian(
     icnf::AbstractICNF{T, <:DIJacVecMatrixMode},
+    ::TestMode,
     f::LuxCore.StatefulLuxLayer,
     xs::AbstractMatrix{<:Real},
-) where {T}
+) where {T <: AbstractFloat}
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(
-        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
-        size(xs, 1),
-        size(xs, 1),
-        size(xs, 2),
-    )
+    res = Zygote.Buffer(y, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[:, i, :] = only(
             DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (z,)),
         )
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
     end
-    return y, eachslice(copy(res); dims = 3)
+    return y, oftype(cat(y; dims = Val(3)), copy(res))
 end
 
-function jacobian_batched(
-    icnf::AbstractICNF{T, <:DIMatrixMode},
+function icnf_jacobian(
+    icnf::AbstractICNF{<:AbstractFloat, <:LuxMatrixMode},
+    ::TestMode,
     f::LuxCore.StatefulLuxLayer,
     xs::AbstractMatrix{<:Real},
-) where {T}
-    y, J = DifferentiationInterface.value_and_jacobian(f, icnf.compute_mode.adback, xs)
-    return y, split_jac(J, size(xs, 1))
+)
+    y = f(xs)
+    return y,
+    oftype(cat(y; dims = Val(3)), Lux.batched_jacobian(f, icnf.compute_mode.adback, xs))
 end
 
-function split_jac(x::AbstractMatrix{<:Real}, sz::Integer)
-    return (
-        x[i:j, i:j] for (i, j) in zip(
-            firstindex(x, 1):sz:lastindex(x, 1),
-            (firstindex(x, 1) + sz - 1):sz:lastindex(x, 1),
-        )
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:DIVecJacVectorMode},
+    ::TrainMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractVector{<:Real},
+    ϵ::AbstractVector{T},
+) where {T <: AbstractFloat}
+    y = f(xs)
+    return y,
+    oftype(
+        y,
+        only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (ϵ,))),
+    )
+end
+
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:DIJacVecVectorMode},
+    ::TrainMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractVector{<:Real},
+    ϵ::AbstractVector{T},
+) where {T <: AbstractFloat}
+    y = f(xs)
+    return y,
+    oftype(
+        y,
+        only(DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (ϵ,))),
+    )
+end
+
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:DIVecJacMatrixMode},
+    ::TrainMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractMatrix{<:Real},
+    ϵ::AbstractMatrix{T},
+) where {T <: AbstractFloat}
+    y = f(xs)
+    return y,
+    oftype(
+        y,
+        only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (ϵ,))),
+    )
+end
+
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:DIJacVecMatrixMode},
+    ::TrainMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractMatrix{<:Real},
+    ϵ::AbstractMatrix{T},
+) where {T <: AbstractFloat}
+    y = f(xs)
+    return y,
+    oftype(
+        y,
+        only(DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (ϵ,))),
     )
 end
 
-function jacobian_batched(
-    icnf::AbstractICNF{T, <:LuxMatrixMode},
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:LuxVecJacMatrixMode},
+    ::TrainMode,
+    f::LuxCore.StatefulLuxLayer,
+    xs::AbstractMatrix{<:Real},
+    ϵ::AbstractMatrix{T},
+) where {T <: AbstractFloat}
+    y = f(xs)
+    return y, oftype(y, Lux.vector_jacobian_product(f, icnf.compute_mode.adback, xs, ϵ))
+end
+
+function icnf_jacobian(
+    icnf::AbstractICNF{T, <:LuxJacVecMatrixMode},
+    ::TrainMode,
     f::LuxCore.StatefulLuxLayer,
     xs::AbstractMatrix{<:Real},
-) where {T}
+    ϵ::AbstractMatrix{T},
+) where {T <: AbstractFloat}
     y = f(xs)
-    J = Lux.batched_jacobian(f, icnf.compute_mode.adback, xs)
-    return y, eachslice(J; dims = 3)
+    return y, oftype(y, Lux.jacobian_vector_product(f, icnf.compute_mode.adback, xs, ϵ))
 end