fix: incorrect activation usage

LuxDL · Jul 23, 2024 · d179933 · d179933 · avik-pal · Jul 23, 2024
1 parent e8f8484
commit d179933
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/src/impl/activation.jl b/src/impl/activation.jl
@@ -34,16 +34,23 @@ end
 _fast_activation(::typeof(identity), x::AbstractArray) = x
 
 @stable default_mode="disable" function _fast_activation(σ::F, x::AbstractArray) where {F}
-    y = similar(x, Core.Compiler._return_type(σ, Tuple{eltype(x)}))
-    _fast_activation!(internal_operation_mode(x), y, σ, x)
-    return y
+    return _fast_activation(internal_operation_mode(x), σ, x)
 end
 
 function CRC.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(_fast_activation),
         σ::F, x::AbstractArray{T}) where {F, T}
     return CRC.rrule_via_ad(cfg, broadcast, σ, x)
 end
 
+_fast_activation(opmode, σ::F, x::AbstractArray) where {F} = broadcast(σ, x)
+
+function _fast_activation(opmode::LoopedArrayOp, σ::F, x::AbstractArray) where {F}
+    RT = Core.Compiler._return_type(σ, Tuple{eltype(x)})
+    y = similar(x, ifelse(isconcretetype(RT), RT, eltype(x)))
+    _fast_activation!(opmode, y, σ, x)
+    return y
+end
+
 _fast_activation!(::typeof(identity), x::AbstractArray) = nothing
 
 @stable default_mode="disable" function _fast_activation!(σ::F, x::AbstractArray) where {F}

diff --git a/src/impl/bias_activation.jl b/src/impl/bias_activation.jl
@@ -142,7 +142,7 @@ function __apply_bias_activation_cached!!(
     if can_setindex(x)
         opmode = internal_operation_mode((x, bias))
         if opmode isa LoopedArrayOp
-            bc = Broadcast.instantiate(Broadcast.broadcasted(σ ∘ +, x, bias_))
+            bc = Broadcast.instantiate(Broadcast.broadcasted(+, x, bias_))
             @simd ivdep for I in eachindex(bc)
                 @inbounds x[I] = bc[I]
             end