fix planar error

impICNF · May 14, 2024 · fb7b747 · fb7b747
1 parent 355fa2c
commit fb7b747
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 11 deletions.
diff --git a/src/ContinuousNormalizingFlows.jl b/src/ContinuousNormalizingFlows.jl
@@ -58,9 +58,10 @@ export construct,
     PlanarLayer,
     MulLayer
 
+include(joinpath("layers", "base_layer.jl"))
 include(joinpath("layers", "cond_layer.jl"))
-include(joinpath("layers", "planar_layer.jl"))
 include(joinpath("layers", "mul_layer.jl"))
+include(joinpath("layers", "planar_layer.jl"))
 
 include("types.jl")
 

diff --git a/src/layers/base_layer.jl b/src/layers/base_layer.jl
@@ -0,0 +1,11 @@
+@inline function apply_act(::typeof(identity), x::Any)
+    x
+end
+
+@inline function apply_act(activation::Any, x::Number)
+    activation(x)
+end
+
+@inline function apply_act(activation::Any, x::AbstractArray)
+    activation.(x)
+end
diff --git a/src/layers/mul_layer.jl b/src/layers/mul_layer.jl
@@ -31,5 +31,5 @@ function LuxCore.outputsize(m::MulLayer)
 end
 
 @inline function (m::MulLayer)(x::AbstractVecOrMat, ps::Any, st::NamedTuple)
-    return Lux.apply_activation(m.activation, Octavian.matmul(ps.weight, x)), st
+    apply_act(m.activation, Octavian.matmul(ps.weight, x)), st
 end
diff --git a/src/layers/planar_layer.jl b/src/layers/planar_layer.jl
@@ -71,34 +71,33 @@ function LuxCore.outputsize(m::PlanarLayer)
 end
 
 @inline function (m::PlanarLayer{true})(z::AbstractVector, ps::Any, st::NamedTuple)
-    ps.u * Lux.apply_bias_activation(m.activation, LinearAlgebra.dot(ps.w, z), only(ps.b)),
-    st
+    ps.u * apply_act(m.activation, LinearAlgebra.dot(ps.w, z) + only(ps.b)), st
 end
 
 @inline function (m::PlanarLayer{true})(z::AbstractMatrix, ps::Any, st::NamedTuple)
-    ps.u * Lux.apply_bias_activation(m.activation, transpose(ps.w) * z, only(ps.b)), st
+    ps.u * apply_act(m.activation, muladd(transpose(ps.w), z, only(ps.b))), st
 end
 
 @inline function (m::PlanarLayer{false})(z::AbstractVector, ps::Any, st::NamedTuple)
-    ps.u * Lux.apply_activation(m.activation, LinearAlgebra.dot(ps.w, z)), st
+    ps.u * apply_act(m.activation, LinearAlgebra.dot(ps.w, z)), st
 end
 
 @inline function (m::PlanarLayer{false})(z::AbstractMatrix, ps::Any, st::NamedTuple)
-    ps.u * Lux.apply_activation(m.activation, transpose(ps.w) * z), st
+    ps.u * apply_act(m.activation, transpose(ps.w) * z), st
 end
 
 @inline function pl_h(m::PlanarLayer{true}, z::AbstractVector, ps::Any, st::NamedTuple)
-    Lux.apply_bias_activation(m.activation, LinearAlgebra.dot(ps.w, z), only(ps.b)), st
+    apply_act(m.activation, LinearAlgebra.dot(ps.w, z) + only(ps.b)), st
 end
 
 @inline function pl_h(m::PlanarLayer{true}, z::AbstractMatrix, ps::Any, st::NamedTuple)
-    Lux.apply_bias_activation(m.activation, transpose(ps.w) * z, only(ps.b)), st
+    apply_act(m.activation, muladd(transpose(ps.w), z, only(ps.b))), st
 end
 
 @inline function pl_h(m::PlanarLayer{false}, z::AbstractVector, ps::Any, st::NamedTuple)
-    Lux.apply_activation(m.activation, LinearAlgebra.dot(ps.w, z)), st
+    apply_act(m.activation, LinearAlgebra.dot(ps.w, z)), st
 end
 
 @inline function pl_h(m::PlanarLayer{false}, z::AbstractMatrix, ps::Any, st::NamedTuple)
-    Lux.apply_activation(m.activation, transpose(ps.w) * z), st
+    apply_act(m.activation, transpose(ps.w) * z), st
 end