direct usage of StatefulLuxLayer (#413)

* direct usage of `StatefulLuxLayer` * fix
impICNF · May 6, 2024 · 0e989bb · 0e989bb
1 parent d6199f8
commit 0e989bb
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 122 deletions.
diff --git a/src/base_icnf.jl b/src/base_icnf.jl
@@ -236,9 +236,9 @@ function inference_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(icnf.nn, ps, st)
+    nn = icnf.nn
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(xs, zrs),
         steer_tspan(icnf, mode),
         ps,
@@ -259,9 +259,9 @@ function inference_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(CondLayer(icnf.nn, ys), ps, st)
+    nn = CondLayer(icnf.nn, ys)
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(xs, zrs),
         steer_tspan(icnf, mode),
         ps,
@@ -281,9 +281,9 @@ function inference_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input, size(xs, 2))
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(icnf.nn, ps, st)
+    nn = icnf.nn
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(xs, zrs),
         steer_tspan(icnf, mode),
         ps,
@@ -304,9 +304,9 @@ function inference_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input, size(xs, 2))
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(CondLayer(icnf.nn, ys), ps, st)
+    nn = CondLayer(icnf.nn, ys)
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(xs, zrs),
         steer_tspan(icnf, mode),
         ps,
@@ -327,9 +327,9 @@ function generate_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(icnf.nn, ps, st)
+    nn = icnf.nn
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(new_xs, zrs),
         reverse(steer_tspan(icnf, mode)),
         ps,
@@ -351,9 +351,9 @@ function generate_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(CondLayer(icnf.nn, ys), ps, st)
+    nn = CondLayer(icnf.nn, ys)
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(new_xs, zrs),
         reverse(steer_tspan(icnf, mode)),
         ps,
@@ -375,9 +375,9 @@ function generate_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input, n)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(icnf.nn, ps, st)
+    nn = icnf.nn
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(new_xs, zrs),
         reverse(steer_tspan(icnf, mode)),
         ps,
@@ -400,9 +400,9 @@ function generate_prob(
     ChainRulesCore.@ignore_derivatives fill!(zrs, zero(T))
     ϵ = base_AT(icnf.resource, icnf, icnf.nvars + n_aug_input, n)
     Random.rand!(icnf.rng, icnf.epsdist, ϵ)
-    nn = Lux.StatefulLuxLayer(CondLayer(icnf.nn, ys), ps, st)
+    nn = CondLayer(icnf.nn, ys)
     SciMLBase.ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
-        make_ode_func(icnf, mode, nn, ϵ),
+        make_ode_func(icnf, mode, nn, st, ϵ),
         vcat(new_xs, zrs),
         reverse(steer_tspan(icnf, mode)),
         ps,
@@ -515,28 +515,21 @@ end
 @inline function make_ode_func(
     icnf::AbstractICNF{T, CM, INPLACE},
     mode::Mode,
-    nn::Lux.StatefulLuxLayer,
+    nn::LuxCore.AbstractExplicitLayer,
+    st::NamedTuple,
     ϵ::AbstractVecOrMat{T},
 ) where {T <: AbstractFloat, CM, INPLACE}
     function ode_func_op(u, p, t)
-        augmented_f(u, p, t, icnf, mode, nn, ϵ)
+        augmented_f(u, p, t, icnf, mode, nn, st, ϵ)
     end
 
     function ode_func_ip(du, u, p, t)
-        augmented_f(du, u, p, t, icnf, mode, nn, ϵ)
+        augmented_f(du, u, p, t, icnf, mode, nn, st, ϵ)
     end
 
     ifelse(INPLACE, ode_func_ip, ode_func_op)
 end
 
-@inline function make_dyn_func(nn::Lux.StatefulLuxLayer, ps::Any)
-    function dyn_func(x)
-        LuxCore.apply(nn, x, ps)
-    end
-
-    dyn_func
-end
-
 @inline function (icnf::AbstractICNF{T, CM, INPLACE, false})(
     xs::AbstractVecOrMat,
     ps::Any,