change types in models (#267)

impICNF · Aug 11, 2023 · 01f16ad · 01f16ad
1 parent e68469c
commit 01f16ad
Show file tree

Hide file tree

Showing 12 changed files with 161 additions and 77 deletions.
diff --git a/src/base.jl b/src/base.jl
@@ -3,8 +3,8 @@ export construct
 function construct(
     aicnf::Type{<:AbstractFlows},
     nn,
-    nvars::Integer,
-    naugmented::Integer = 0;
+    nvars::Int,
+    naugmented::Int = 0;
     data_type::Type{<:AbstractFloat} = Float32,
     compute_mode::Type{<:ComputeMode} = ADVectorMode,
     resource::AbstractResource = CPU1(),
@@ -23,7 +23,17 @@ function construct(
 )
     steerdist = Uniform{data_type}(-steer_rate, steer_rate)
 
-    aicnf{data_type, compute_mode, !iszero(naugmented), !iszero(steer_rate)}(
+    aicnf{
+        data_type,
+        compute_mode,
+        !iszero(naugmented),
+        !iszero(steer_rate),
+        typeof(resource),
+        typeof(basedist),
+        typeof(tspan),
+        typeof(steerdist),
+        typeof(differentiation_backend),
+    }(
         nn,
         nvars,
         naugmented,

diff --git a/src/base_cond_icnf.jl b/src/base_cond_icnf.jl
@@ -263,7 +263,7 @@ function generate_prob(
     ys::AbstractMatrix{<:Real},
     ps::Any,
     st::Any,
-    n::Integer;
+    n::Int;
     resource::AbstractResource = icnf.resource,
     tspan::NTuple{2} = icnf.tspan,
     steerdist::Distribution = icnf.steerdist,
@@ -310,7 +310,7 @@ function generate(
     ys::AbstractMatrix{<:Real},
     ps::Any,
     st::Any,
-    n::Integer;
+    n::Int;
     resource::AbstractResource = icnf.resource,
     tspan::NTuple{2} = icnf.tspan,
     steerdist::Distribution = icnf.steerdist,

diff --git a/src/base_icnf.jl b/src/base_icnf.jl
@@ -249,7 +249,7 @@ function generate_prob(
     mode::Mode,
     ps::Any,
     st::Any,
-    n::Integer;
+    n::Int;
     resource::AbstractResource = icnf.resource,
     tspan::NTuple{2} = icnf.tspan,
     steerdist::Distribution = icnf.steerdist,
@@ -294,7 +294,7 @@ function generate(
     mode::Mode,
     ps::Any,
     st::Any,
-    n::Integer;
+    n::Int;
     resource::AbstractResource = icnf.resource,
     tspan::NTuple{2} = icnf.tspan,
     steerdist::Distribution = icnf.steerdist,

diff --git a/src/cond_ffjord.jl b/src/cond_ffjord.jl
@@ -3,17 +3,26 @@ export CondFFJORD
 """
 Implementation of FFJORD (Conditional Version)
 """
-struct CondFFJORD{T <: AbstractFloat, CM <: ComputeMode, AUGMENTED, STEER} <:
-       AbstractCondICNF{T, CM, AUGMENTED, STEER}
+struct CondFFJORD{
+    T <: AbstractFloat,
+    CM <: ComputeMode,
+    AUGMENTED,
+    STEER,
+    RESOURCE <: AbstractResource,
+    BASEDIST <: Distribution,
+    TSPAN <: NTuple{2, T},
+    STEERDIST <: Distribution,
+    DIFFERENTIATION_BACKEND <: AbstractDifferentiation.AbstractBackend,
+} <: AbstractCondICNF{T, CM, AUGMENTED, STEER}
     nn::LuxCore.AbstractExplicitLayer
-    nvars::Integer
-    naugmented::Integer
+    nvars::Int
+    naugmented::Int
 
-    resource::AbstractResource
-    basedist::Distribution
-    tspan::NTuple{2, T}
-    steerdist::Distribution
-    differentiation_backend::AbstractDifferentiation.AbstractBackend
+    resource::RESOURCE
+    basedist::BASEDIST
+    tspan::TSPAN
+    steerdist::STEERDIST
+    differentiation_backend::DIFFERENTIATION_BACKEND
     sol_args::Tuple
     sol_kwargs::Dict
 end

diff --git a/src/cond_planar.jl b/src/cond_planar.jl
@@ -3,17 +3,26 @@ export CondPlanar
 """
 Implementation of Planar (Conditional Version)
 """
-struct CondPlanar{T <: AbstractFloat, CM <: ComputeMode, AUGMENTED, STEER} <:
-       AbstractCondICNF{T, CM, AUGMENTED, STEER}
+struct CondPlanar{
+    T <: AbstractFloat,
+    CM <: ComputeMode,
+    AUGMENTED,
+    STEER,
+    RESOURCE <: AbstractResource,
+    BASEDIST <: Distribution,
+    TSPAN <: NTuple{2, T},
+    STEERDIST <: Distribution,
+    DIFFERENTIATION_BACKEND <: AbstractDifferentiation.AbstractBackend,
+} <: AbstractCondICNF{T, CM, AUGMENTED, STEER}
     nn::PlanarLayer
-    nvars::Integer
-    naugmented::Integer
+    nvars::Int
+    naugmented::Int
 
-    resource::AbstractResource
-    basedist::Distribution
-    tspan::NTuple{2, T}
-    steerdist::Distribution
-    differentiation_backend::AbstractDifferentiation.AbstractBackend
+    resource::RESOURCE
+    basedist::BASEDIST
+    tspan::TSPAN
+    steerdist::STEERDIST
+    differentiation_backend::DIFFERENTIATION_BACKEND
     sol_args::Tuple
     sol_kwargs::Dict
 end

diff --git a/src/cond_rnode.jl b/src/cond_rnode.jl
@@ -3,17 +3,26 @@ export CondRNODE
 """
 Implementation of RNODE (Conditional Version)
 """
-struct CondRNODE{T <: AbstractFloat, CM <: ComputeMode, AUGMENTED, STEER} <:
-       AbstractCondICNF{T, CM, AUGMENTED, STEER}
+struct CondRNODE{
+    T <: AbstractFloat,
+    CM <: ComputeMode,
+    AUGMENTED,
+    STEER,
+    RESOURCE <: AbstractResource,
+    BASEDIST <: Distribution,
+    TSPAN <: NTuple{2, T},
+    STEERDIST <: Distribution,
+    DIFFERENTIATION_BACKEND <: AbstractDifferentiation.AbstractBackend,
+} <: AbstractCondICNF{T, CM, AUGMENTED, STEER}
     nn::LuxCore.AbstractExplicitLayer
-    nvars::Integer
-    naugmented::Integer
+    nvars::Int
+    naugmented::Int
 
-    resource::AbstractResource
-    basedist::Distribution
-    tspan::NTuple{2, T}
-    steerdist::Distribution
-    differentiation_backend::AbstractDifferentiation.AbstractBackend
+    resource::RESOURCE
+    basedist::BASEDIST
+    tspan::TSPAN
+    steerdist::STEERDIST
+    differentiation_backend::DIFFERENTIATION_BACKEND
     sol_args::Tuple
     sol_kwargs::Dict
     λ₁::T
@@ -23,8 +32,8 @@ end
 function construct(
     aicnf::Type{<:CondRNODE},
     nn,
-    nvars::Integer,
-    naugmented::Integer = 0;
+    nvars::Int,
+    naugmented::Int = 0;
     data_type::Type{<:AbstractFloat} = Float32,
     compute_mode::Type{<:ComputeMode} = ADVectorMode,
     resource::AbstractResource = CPU1(),
@@ -45,7 +54,17 @@ function construct(
 )
     steerdist = Uniform{data_type}(-steer_rate, steer_rate)
 
-    aicnf{data_type, compute_mode, !iszero(naugmented), !iszero(steer_rate)}(
+    aicnf{
+        data_type,
+        compute_mode,
+        !iszero(naugmented),
+        !iszero(steer_rate),
+        typeof(resource),
+        typeof(basedist),
+        typeof(tspan),
+        typeof(steerdist),
+        typeof(differentiation_backend),
+    }(
         nn,
         nvars,
         naugmented,

diff --git a/src/core_cond_icnf.jl b/src/core_cond_icnf.jl
@@ -7,11 +7,11 @@ mutable struct CondICNFModel <: MLJICNF
     loss::Function
 
     optimizers::AbstractVector
-    n_epochs::Integer
+    n_epochs::Int
     adtype::ADTypes.AbstractADType
 
     use_batch::Bool
-    batch_size::Integer
+    batch_size::Int
     have_callback::Bool
 
     data_type::Type{<:AbstractFloat}
@@ -22,10 +22,10 @@ function CondICNFModel(
     m::AbstractCondICNF{T, CM},
     loss::Function = loss;
     optimizers::AbstractVector = Any[Optimisers.Lion(),],
-    n_epochs::Integer = 300,
+    n_epochs::Int = 300,
     adtype::ADTypes.AbstractADType = AutoZygote(),
     use_batch::Bool = true,
-    batch_size::Integer = 32,
+    batch_size::Int = 32,
     have_callback::Bool = true,
 ) where {T <: AbstractFloat, CM <: ComputeMode}
     CondICNFModel(

diff --git a/src/core_icnf.jl b/src/core_icnf.jl
@@ -7,11 +7,11 @@ mutable struct ICNFModel <: MLJICNF
     loss::Function
 
     optimizers::AbstractVector
-    n_epochs::Integer
+    n_epochs::Int
     adtype::ADTypes.AbstractADType
 
     use_batch::Bool
-    batch_size::Integer
+    batch_size::Int
     have_callback::Bool
 
     data_type::Type{<:AbstractFloat}
@@ -22,10 +22,10 @@ function ICNFModel(
     m::AbstractICNF{T, CM},
     loss::Function = loss;
     optimizers::AbstractVector = Any[Optimisers.Lion(),],
-    n_epochs::Integer = 300,
+    n_epochs::Int = 300,
     adtype::ADTypes.AbstractADType = AutoZygote(),
     use_batch::Bool = true,
-    batch_size::Integer = 32,
+    batch_size::Int = 32,
     have_callback::Bool = true,
 ) where {T <: AbstractFloat, CM <: ComputeMode}
     ICNFModel(

diff --git a/src/ffjord.jl b/src/ffjord.jl
@@ -5,17 +5,26 @@ Implementation of FFJORD from
 
 [Grathwohl, Will, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. "Ffjord: Free-form continuous dynamics for scalable reversible generative models." arXiv preprint arXiv:1810.01367 (2018).](https://arxiv.org/abs/1810.01367)
 """
-struct FFJORD{T <: AbstractFloat, CM <: ComputeMode, AUGMENTED, STEER} <:
-       AbstractICNF{T, CM, AUGMENTED, STEER}
+struct FFJORD{
+    T <: AbstractFloat,
+    CM <: ComputeMode,
+    AUGMENTED,
+    STEER,
+    RESOURCE <: AbstractResource,
+    BASEDIST <: Distribution,
+    TSPAN <: NTuple{2, T},
+    STEERDIST <: Distribution,
+    DIFFERENTIATION_BACKEND <: AbstractDifferentiation.AbstractBackend,
+} <: AbstractICNF{T, CM, AUGMENTED, STEER}
     nn::LuxCore.AbstractExplicitLayer
-    nvars::Integer
-    naugmented::Integer
+    nvars::Int
+    naugmented::Int
 
-    resource::AbstractResource
-    basedist::Distribution
-    tspan::NTuple{2, T}
-    steerdist::Distribution
-    differentiation_backend::AbstractDifferentiation.AbstractBackend
+    resource::RESOURCE
+    basedist::BASEDIST
+    tspan::TSPAN
+    steerdist::STEERDIST
+    differentiation_backend::DIFFERENTIATION_BACKEND
     sol_args::Tuple
     sol_kwargs::Dict
 end

diff --git a/src/planar.jl b/src/planar.jl
@@ -5,17 +5,26 @@ Implementation of Planar Flows from
 
 [Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. "Neural Ordinary Differential Equations." arXiv preprint arXiv:1806.07366 (2018).](https://arxiv.org/abs/1806.07366)
 """
-struct Planar{T <: AbstractFloat, CM <: ComputeMode, AUGMENTED, STEER} <:
-       AbstractICNF{T, CM, AUGMENTED, STEER}
+struct Planar{
+    T <: AbstractFloat,
+    CM <: ComputeMode,
+    AUGMENTED,
+    STEER,
+    RESOURCE <: AbstractResource,
+    BASEDIST <: Distribution,
+    TSPAN <: NTuple{2, T},
+    STEERDIST <: Distribution,
+    DIFFERENTIATION_BACKEND <: AbstractDifferentiation.AbstractBackend,
+} <: AbstractICNF{T, CM, AUGMENTED, STEER}
     nn::PlanarLayer
-    nvars::Integer
-    naugmented::Integer
+    nvars::Int
+    naugmented::Int
 
-    resource::AbstractResource
-    basedist::Distribution
-    tspan::NTuple{2, T}
-    steerdist::Distribution
-    differentiation_backend::AbstractDifferentiation.AbstractBackend
+    resource::RESOURCE
+    basedist::BASEDIST
+    tspan::TSPAN
+    steerdist::STEERDIST
+    differentiation_backend::DIFFERENTIATION_BACKEND
     sol_args::Tuple
     sol_kwargs::Dict
 end

diff --git a/src/planar_layer.jl b/src/planar_layer.jl
@@ -10,7 +10,7 @@ struct PlanarLayer{use_bias, cond, F1, F2, F3} <: LuxCore.AbstractExplicitLayer
     nvars::Int
     init_weight::F2
     init_bias::F3
-    n_cond::Integer
+    n_cond::Int
 end
 
 function PlanarLayer(
@@ -21,7 +21,7 @@ function PlanarLayer(
     use_bias::Bool = true,
     allow_fast_activation::Bool = true,
     cond::Bool = false,
-    n_cond::Integer = 0,
+    n_cond::Int = 0,
 )
     activation = allow_fast_activation ? NNlib.fast_act(activation) : activation
     PlanarLayer{use_bias, cond, typeof(activation), typeof(init_weight), typeof(init_bias)}(