Better TestMode

Project.toml

@@ -7,6 +7,7 @@ version = "0.6.0"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ComputationalResources = "ed09eef8-17a6-5b46-8889-db040fac31e3"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
@@ -44,6 +45,7 @@ cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 ADTypes = "0.1"
 AbstractDifferentiation = "0.5"
 CUDA = "4"
+ChainRulesCore = "1"
 ComponentArrays = "0.13"
 ComputationalResources = "0.3"
 DataFrames = "1"

README.md

@@ -63,7 +63,7 @@ fit!(mach)
 ps, st = fitted_params(mach)
 
 # Use It
-d = ICNFDist(icnf, ps, st)
+d = ICNFDist(icnf, mach, TestMode())
 actual_pdf = pdf.(data_dist, vec(r))
 estimated_pdf = pdf(d, r)
 new_data = rand(d, n)

src/ContinuousNormalizingFlows.jl

@@ -3,6 +3,7 @@ module ContinuousNormalizingFlows
 using AbstractDifferentiation,
     ADTypes,
     Base.Iterators,
+    ChainRulesCore,
     ComponentArrays,
     ComputationalResources,
     CUDA,

src/base.jl

@@ -3,8 +3,7 @@ export construct
 function construct(
     aicnf::Type{<:AbstractFlows},
     nn,
-    nvars::Integer,
-    ;
+    nvars::Integer;
     data_type::Type{<:AbstractFloat} = Float32,
     array_type::Type{<:AbstractArray} = Array,
     compute_mode::Type{<:ComputeMode} = ADVectorMode,

src/base_cond_icnf.jl

@@ -111,12 +111,12 @@ end
 
 function loss(
     icnf::AbstractCondICNF{T, AT, <:VectorMode},
+    mode::Mode,
     xs::AbstractVector{<:Real},
     ys::AbstractVector{<:Real},
     ps::Any,
     st::Any;
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,
@@ -138,12 +138,12 @@ end
 
 function loss(
     icnf::AbstractCondICNF{T, AT, <:MatrixMode},
+    mode::Mode,
     xs::AbstractMatrix{<:Real},
     ys::AbstractMatrix{<:Real},
     ps::Any,
     st::Any;
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,

src/base_icnf.jl

@@ -107,11 +107,11 @@ end
 
 function loss(
     icnf::AbstractICNF{T, AT, <:VectorMode},
+    mode::Mode,
     xs::AbstractVector{<:Real},
     ps::Any,
     st::Any;
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,
@@ -132,11 +132,11 @@ end
 
 function loss(
     icnf::AbstractICNF{T, AT, <:MatrixMode},
+    mode::Mode,
     xs::AbstractMatrix{<:Real},
     ps::Any,
     st::Any;
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,

src/cond_rnode.jl

@@ -178,14 +178,14 @@ end
 
 function loss(
     icnf::CondRNODE{T, AT, <:VectorMode},
+    mode::TrainMode,
     xs::AbstractVector{<:Real},
     ys::AbstractVector{<:Real},
     ps::Any,
     st::Any,
     λ₁::T = convert(T, 1e-2),
     λ₂::T = convert(T, 1e-2);
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,
@@ -207,14 +207,14 @@ end
 
 function loss(
     icnf::CondRNODE{T, AT, <:MatrixMode},
+    mode::TrainMode,
     xs::AbstractMatrix{<:Real},
     ys::AbstractMatrix{<:Real},
     ps::Any,
     st::Any,
     λ₁::T = convert(T, 1e-2),
     λ₂::T = convert(T, 1e-2);
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,

src/core_cond_icnf.jl

@@ -4,7 +4,7 @@ export CondICNFModel, CondICNFDist
 
 function loss_f(icnf::AbstractCondICNF, loss::Function, st::Any)::Function
     function f(ps, θ, xs, ys)
-        loss(icnf, xs, ys, ps, st)
+        loss(icnf, TrainMode(), xs, ys, ps, st)
     end
     f
 end
@@ -31,8 +31,7 @@ end
 
 function CondICNFModel(
     m::AbstractCondICNF{T, AT, CM},
-    loss::Function = loss,
-    ;
+    loss::Function = loss;
     optimizers::AbstractVector = [Optimisers.Adam()],
     n_epochs::Integer = 300,
     adtype::ADTypes.AbstractADType = AutoZygote(),
@@ -192,16 +191,26 @@ MLJBase.metadata_model(
 
 struct CondICNFDist <: ICNFDistribution
     m::AbstractCondICNF
+    mode::Mode
     ys::AbstractVecOrMat{<:Real}
     ps::Any
     st::Any
 end
 
+function CondICNFDist(
+    mach::Machine{<:CondICNFModel},
+    mode::Mode,
+    ys::AbstractVecOrMat{<:Real},
+)
+    (ps, st) = fitted_params(mach)
+    CondICNFDist(mach.model.m, mode, ys, ps, st)
+end
+
 Base.length(d::CondICNFDist) = d.m.nvars
 Base.eltype(d::CondICNFDist) = typeof(d.m).parameters[1]
 function Distributions._logpdf(d::CondICNFDist, x::AbstractVector{<:Real})
     if d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
-        first(inference(d.m, TestMode(), x, d.ys, d.ps, d.st))
+        first(inference(d.m, d.mode, x, d.ys, d.ps, d.st))
     elseif d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
         first(Distributions._logpdf(d, hcat(x)))
     else
@@ -212,14 +221,14 @@ function Distributions._logpdf(d::CondICNFDist, A::AbstractMatrix{<:Real})
     if d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
         broadcast(x -> Distributions._logpdf(d, x), eachcol(A))
     elseif d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
-        first(inference(d.m, TestMode(), A, d.ys[:, begin:size(A, 2)], d.ps, d.st))
+        first(inference(d.m, d.mode, A, d.ys[:, begin:size(A, 2)], d.ps, d.st))
     else
         error("Not Implemented")
     end
 end
 function Distributions._rand!(rng::AbstractRNG, d::CondICNFDist, x::AbstractVector{<:Real})
     if d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
-        x .= generate(d.m, TestMode(), d.ys, d.ps, d.st; rng)
+        x .= generate(d.m, d.mode, d.ys, d.ps, d.st; rng)
     elseif d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
         x .= Distributions._rand!(rng, d, hcat(x))
     else
@@ -230,15 +239,7 @@ function Distributions._rand!(rng::AbstractRNG, d::CondICNFDist, A::AbstractMatr
     if d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
         A .= hcat(broadcast(x -> Distributions._rand!(rng, d, x), eachcol(A))...)
     elseif d.m isa AbstractCondICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
-        A .= generate(
-            d.m,
-            TestMode(),
-            d.ys[:, begin:size(A, 2)],
-            d.ps,
-            d.st,
-            size(A, 2);
-            rng,
-        )
+        A .= generate(d.m, d.mode, d.ys[:, begin:size(A, 2)], d.ps, d.st, size(A, 2); rng)
     else
         error("Not Implemented")
     end

src/core_icnf.jl

@@ -4,7 +4,7 @@ export ICNFModel, ICNFDist
 
 function loss_f(icnf::AbstractICNF, loss::Function, st::Any)::Function
     function f(ps, θ, xs)
-        loss(icnf, xs, ps, st)
+        loss(icnf, TrainMode(), xs, ps, st)
     end
     f
 end
@@ -31,8 +31,7 @@ end
 
 function ICNFModel(
     m::AbstractICNF{T, AT, CM},
-    loss::Function = loss,
-    ;
+    loss::Function = loss;
     optimizers::AbstractVector = [Optimisers.Adam()],
     n_epochs::Integer = 300,
     adtype::ADTypes.AbstractADType = AutoZygote(),
@@ -183,15 +182,21 @@ MLJBase.metadata_model(
 
 struct ICNFDist <: ICNFDistribution
     m::AbstractICNF
+    mode::Mode
     ps::Any
     st::Any
 end
 
+function ICNFDist(mach::Machine{<:ICNFModel}, mode::Mode)
+    (ps, st) = fitted_params(mach)
+    ICNFDist(mach.model.m, mode, ps, st)
+end
+
 Base.length(d::ICNFDist) = d.m.nvars
 Base.eltype(d::ICNFDist) = typeof(d.m).parameters[1]
 function Distributions._logpdf(d::ICNFDist, x::AbstractVector{<:Real})
     if d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
-        first(inference(d.m, TestMode(), x, d.ps, d.st))
+        first(inference(d.m, d.mode, x, d.ps, d.st))
     elseif d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
         first(Distributions._logpdf(d, hcat(x)))
     else
@@ -202,14 +207,14 @@ function Distributions._logpdf(d::ICNFDist, A::AbstractMatrix{<:Real})
     if d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
         broadcast(x -> Distributions._logpdf(d, x), eachcol(A))
     elseif d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
-        first(inference(d.m, TestMode(), A, d.ps, d.st))
+        first(inference(d.m, d.mode, A, d.ps, d.st))
     else
         error("Not Implemented")
     end
 end
 function Distributions._rand!(rng::AbstractRNG, d::ICNFDist, x::AbstractVector{<:Real})
     if d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
-        x .= generate(d.m, TestMode(), d.ps, d.st; rng)
+        x .= generate(d.m, d.mode, d.ps, d.st; rng)
     elseif d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
         x .= Distributions._rand!(rng, d, hcat(x))
     else
@@ -220,7 +225,7 @@ function Distributions._rand!(rng::AbstractRNG, d::ICNFDist, A::AbstractMatrix{<
     if d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:VectorMode}
         A .= hcat(broadcast(x -> Distributions._rand!(rng, d, x), eachcol(A))...)
     elseif d.m isa AbstractICNF{<:AbstractFloat, <:AbstractArray, <:MatrixMode}
-        A .= generate(d.m, TestMode(), d.ps, d.st, size(A, 2); rng)
+        A .= generate(d.m, d.mode, d.ps, d.st, size(A, 2); rng)
     else
         error("Not Implemented")
     end

src/precompile.jl

@@ -1,16 +1,17 @@
 @setup_workload begin
     @compile_workload begin
         rng = Random.default_rng()
-        mts = Type{<:AbstractICNF}[RNODE, FFJORD, Planar]
-        cmts = Type{<:AbstractCondICNF}[CondRNODE, CondFFJORD, CondPlanar]
+        mts = Type{<:AbstractICNF}[RNODE]
+        cmts = Type{<:AbstractCondICNF}[CondRNODE]
         cmodes = Type{<:ComputeMode}[ADVectorMode, ZygoteMatrixMode, SDVecJacMatrixMode]
+        omodes = Mode[TrainMode(), TestMode()]
         nvars = 2
         r = rand(Float32, nvars)
         r_arr = rand(Float32, nvars, 2)
         r2 = rand(Float32, nvars)
         r2_arr = rand(Float32, nvars, 2)
 
-        for cmode in cmodes, mt in mts
+        for cmode in cmodes, omode in omodes, mt in mts
             if mt <: Planar
                 nn = PlanarLayer(nvars, tanh)
             else
@@ -20,12 +21,12 @@
             ps, st = Lux.setup(rng, icnf)
             ps = ComponentArray(ps)
             if cmode <: VectorMode
-                L = loss(icnf, r, ps, st)
+                L = loss(icnf, omode, r, ps, st)
             elseif cmode <: MatrixMode
-                L = loss(icnf, r_arr, ps, st)
+                L = loss(icnf, omode, r_arr, ps, st)
             end
         end
-        for cmode in cmodes, mt in cmts
+        for cmode in cmodes, omode in omodes, mt in cmts
             if mt <: CondPlanar
                 nn = PlanarLayer(nvars, tanh; cond = true)
             else
@@ -35,9 +36,9 @@
             ps, st = Lux.setup(rng, icnf)
             ps = ComponentArray(ps)
             if cmode <: VectorMode
-                L = loss(icnf, r, r2, ps, st)
+                L = loss(icnf, omode, r, r2, ps, st)
             elseif cmode <: MatrixMode
-                L = loss(icnf, r_arr, r2_arr, ps, st)
+                L = loss(icnf, omode, r_arr, r2_arr, ps, st)
             end
         end
     end

src/rnode.jl

@@ -167,13 +167,13 @@ end
 
 function loss(
     icnf::RNODE{T, AT, <:VectorMode},
+    mode::TrainMode,
     xs::AbstractVector{<:Real},
     ps::Any,
     st::Any,
     λ₁::T = convert(T, 1e-2),
     λ₂::T = convert(T, 1e-2);
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,
@@ -194,13 +194,13 @@ end
 
 function loss(
     icnf::RNODE{T, AT, <:MatrixMode},
+    mode::TrainMode,
     xs::AbstractMatrix{<:Real},
     ps::Any,
     st::Any,
     λ₁::T = convert(T, 1e-2),
     λ₂::T = convert(T, 1e-2);
     differentiation_backend::AbstractDifferentiation.AbstractBackend = icnf.differentiation_backend,
-    mode::Mode = TrainMode(),
     rng::AbstractRNG = Random.default_rng(),
     sol_args::Tuple = icnf.sol_args,
     sol_kwargs::Dict = icnf.sol_kwargs,