start fixing type instability

Author: itsdfish

src/conditional_layers/conditional_layer_glow.jl

@@ -58,11 +58,11 @@ or
 
  See also: [`Conv1x1`](@ref), [`ResidualBlock`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-struct ConditionalLayerGlow <: NeuralNetLayer
-    C::Conv1x1
+struct ConditionalLayerGlow{AF,Conv} <: NeuralNetLayer
+    C::Conv
     RB::ResidualBlock
     logdet::Bool
-    activation::ActivationFunction
+    activation::AF
 end
 
 @Flux.functor ConditionalLayerGlow

src/layers/invertible_layer_actnorm.jl

@@ -39,10 +39,10 @@ export ActNorm, reset!
 
  See also: [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-mutable struct ActNorm <: NeuralNetLayer
-    k::Integer
-    s::Parameter
-    b::Parameter
+mutable struct ActNorm{P1,P2,I<:Integer} <: NeuralNetLayer
+    k::I
+    s::P1
+    b::P2
     logdet::Bool
     is_reversed::Bool
 end

src/layers/invertible_layer_basic.jl

@@ -58,10 +58,10 @@ or
 
  See also: [`ResidualBlock`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-mutable struct CouplingLayerBasic <: NeuralNetLayer
+mutable struct CouplingLayerBasic{AF} <: NeuralNetLayer
     RB::Union{ResidualBlock, FluxBlock}
     logdet::Bool
-    activation::ActivationFunction
+    activation::AF
     is_reversed::Bool
 end
 

src/layers/invertible_layer_conv1x1.jl

@@ -39,11 +39,11 @@ export Conv1x1
 
  See also: [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-struct Conv1x1 <: NeuralNetLayer
-    k::Integer
-    v1::Parameter
-    v2::Parameter
-    v3::Parameter
+struct Conv1x1{P1,P2,P3,I<:Integer} <: NeuralNetLayer
+    k::I
+    v1::P1
+    v2::P2
+    v3::P3
     logdet::Bool
     freeze::Bool
 end

src/layers/invertible_layer_glow.jl

@@ -60,8 +60,8 @@ or
 
  See also: [`Conv1x1`](@ref), [`ResidualBlock`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-struct CouplingLayerGlow <: NeuralNetLayer
-    C::Conv1x1
+struct CouplingLayerGlow{Conv<:Conv1x1} <: NeuralNetLayer
+    C::Conv
     RB::Union{ResidualBlock, FluxBlock}
     logdet::Bool
     activation::ActivationFunction

src/networks/invertible_network_conditional_glow.jl

@@ -59,9 +59,9 @@ export NetworkConditionalGlow, NetworkConditionalGlow3D
 
  See also: [`ActNorm`](@ref), [`CouplingLayerGlow!`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
 """
-struct NetworkConditionalGlow <: InvertibleNetwork
-    AN::AbstractArray{ActNorm, 2}
-    AN_C::ActNorm
+struct NetworkConditionalGlow{AN} <: InvertibleNetwork
+    AN::AbstractArray{AN, 2}
+    AN_C::AN
     CL::AbstractArray{ConditionalLayerGlow, 2}
     Z_dims::Union{Array{Array, 1}, Nothing}
     L::Int64

src/utils/dimensionality_operations.jl

@@ -9,9 +9,9 @@ export ShuffleLayer, WaveletLayer, HaarLayer
 ###############################################################################
 # Custom type for squeezer functions
 
-struct Squeezer
-    forward::Function
-    inverse::Function
+struct Squeezer{F1,F2}
+    forward::F1
+    inverse::F2
 end
 
 function ShuffleLayer(;pattern="checkerboard")