add non-autonomous flow

Author: prbzrg

Project.toml

@@ -12,7 +12,6 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
@@ -42,7 +41,6 @@ DifferentiationInterface = "0.7"
 Distributions = "0.25"
 DistributionsAD = "0.6"
 FillArrays = "1"
-ForwardDiff = "1"
 LinearAlgebra = "1"
 Lux = "1"
 LuxCore = "1"

benchmark/Project.toml

@@ -4,8 +4,6 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -17,8 +15,6 @@ BenchmarkTools = "1"
 ComponentArrays = "0.15"
 DifferentiationInterface = "0.7"
 Distributions = "0.25"
-ForwardDiff = "1"
-Lux = "1"
 LuxCore = "1"
 PkgBenchmark = "0.2"
 StableRNGs = "1"

benchmark/benchmarks.jl

@@ -3,8 +3,6 @@ import ADTypes,
     ComponentArrays,
     DifferentiationInterface,
     Distributions,
-    ForwardDiff,
-    Lux,
     LuxCore,
     PkgBenchmark,
     StableRNGs,
@@ -19,18 +17,8 @@ r = rand(rng, data_dist, ndimension, ndata)
 r = convert.(Float32, r)
 
 nvars = size(r, 1)
-naugs = nvars + 1
-n_in = nvars + naugs
-
-nn = Lux.Chain(
-    Lux.Dense(n_in => (2 * n_in + 1), tanh),
-    Lux.Dense((2 * n_in + 1) => n_in, tanh),
-)
-
-icnf = ContinuousNormalizingFlows.ICNF(; nn, nvars, naugmented = naugs, rng)
-
-icnf2 =
-    ContinuousNormalizingFlows.ICNF(; nn, nvars, naugmented = naugs, rng, inplace = true)
+icnf = ContinuousNormalizingFlows.ICNF(; nvars, rng)
+icnf2 = ContinuousNormalizingFlows.ICNF(; nvars, rng, inplace = true)
 
 ps, st = LuxCore.setup(icnf.rng, icnf)
 ps = ComponentArrays.ComponentArray(ps)

examples/usage.jl

@@ -31,7 +31,7 @@ using ContinuousNormalizingFlows,
 # To use gpu, add related packages
 # using LuxCUDA
 
-nn = Chain(Dense(n_in => (2 * n_in + 1), tanh), Dense((2 * n_in + 1) => n_in, tanh))
+nn = Chain(Dense(n_in + 1 => n_in, tanh))
 icnf = ICNF(;
     nn = nn,
     nvars = nvars, # number of variables
@@ -45,6 +45,7 @@ icnf = ICNF(;
     # device = gpu_device(), # process data by GPU
     cond = false, # not conditioning on auxiliary input
     inplace = false, # not using the inplace version of functions
+    autonomous = false, # using non-autonomous flow
     compute_mode = LuxVecJacMatrixMode(AutoZygote()), # process data in batches and use Zygote
     sol_kwargs = (;
         save_everystep = false,

src/ContinuousNormalizingFlows.jl

@@ -8,7 +8,6 @@ import ADTypes,
     Distributions,
     DistributionsAD,
     FillArrays,
-    ForwardDiff,
     LinearAlgebra,
     Lux,
     LuxCore,