Merge branch 'main' into polyalg-23

impICNF · May 8, 2024 · 0a972e6 · 0a972e6
2 parents 741beb1 + 355fa2c
commit 0a972e6
Show file tree

Hide file tree

Showing 13 changed files with 110 additions and 133 deletions.
diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml
@@ -16,7 +16,7 @@ jobs:
           version: '1'
           arch: x64
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - name: Install dependencies
         shell: julia --color=yes {0}

diff --git a/.github/workflows/CI-ENH.yml b/.github/workflows/CI-ENH.yml
@@ -45,7 +45,7 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         with:

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -49,7 +49,7 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         with:

diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
@@ -16,7 +16,7 @@ jobs:
           version: '1'
           arch: x64
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - name: Pkg.add
         shell: julia --color=yes {0}
         run: |

diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml
@@ -25,7 +25,7 @@ jobs:
           version: '1'
           arch: x64
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - name: Configure doc environment
         shell: julia --project=docs --color=yes {0}
         run: |

diff --git a/.github/workflows/Formatter.yml b/.github/workflows/Formatter.yml
@@ -16,7 +16,7 @@ jobs:
           version: '1'
           arch: x64
           show-versioninfo: true
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - name: Install JuliaFormatter and format
         shell: julia --color=yes {0}
         run: |

diff --git a/.github/workflows/Invalidations.yml b/.github/workflows/Invalidations.yml
@@ -19,7 +19,7 @@ jobs:
         version: '1'
         arch: x64
         show-versioninfo: true
-    - uses: julia-actions/cache@v1
+    - uses: julia-actions/cache@v2
     - uses: julia-actions/julia-buildpkg@v1
     - uses: julia-actions/julia-invalidations@v1
       id: invs_pr

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ContinuousNormalizingFlows"
 uuid = "00b1973d-5b2e-40bf-8604-5c9c1d8f50ac"
 authors = ["Hossein Pourbozorg <[email protected]> and contributors"]
-version = "0.22.1"
+version = "0.22.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -50,7 +50,7 @@ ComputationalResources = "0.3"
 DataFrames = "1"
 Dates = "1"
 DifferentialEquations = "7"
-DifferentiationInterface = "0.1, 0.2"
+DifferentiationInterface = "0.1, 0.2, 0.3"
 Distributions = "0.25"
 DistributionsAD = "0.6"
 FillArrays = "1"

diff --git a/benchmark/Project.toml b/benchmark/Project.toml
@@ -12,7 +12,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ADTypes = "0.2, 1"
 BenchmarkTools = "1"
 ComponentArrays = "0.15"
-DifferentiationInterface = "0.1, 0.2"
+DifferentiationInterface = "0.1, 0.2, 0.3"
 Lux = "0.5"
 PkgBenchmark = "0.2"
 StableRNGs = "1"

diff --git a/src/base_icnf.jl b/src/base_icnf.jl
@@ -233,9 +233,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,
@@ -256,9 +256,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,
@@ -278,9 +278,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,
@@ -301,9 +301,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,
@@ -324,9 +324,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,
@@ -348,9 +348,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,
@@ -372,9 +372,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,
@@ -397,9 +397,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,
@@ -512,28 +512,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,