test: fused bias activation

Author: avik-pal

src/TracedRArray.jl

@@ -89,7 +89,7 @@ end
 
 Base.size(x::TracedRArray) = x.shape
 
-Base.copy(A::TracedRArray{T,N}) where {T,N} = TracedRArray((), A.mlir_data, size(A))
+Base.copy(A::TracedRArray{T,N}) where {T,N} = TracedRArray{T,N}((), A.mlir_data, size(A))
 
 function Base.similar(x::TracedRArray{T,N}, ::Type{T2}) where {T,N,T2}
     return TracedRArray{T2,N}((), nothing, size(x))

test/nn/luxlib.jl

@@ -3,6 +3,76 @@ using LuxLib, Reactant, Enzyme, NNlib
 @testset "Fused Dense" begin end
 
 @testset "Bias Activation" begin
+    biasact(act, x, b) = bias_activation(act, x, b)
+    sumabs2biasact(act, x, b) = sum(abs2, biasact(act, x, b))
+    biasact!!(act, x, b) = bias_activation!!(act, copy(x), b)
+    sumabs2biasact!!(act, x, b) = sum(abs2, biasact!!(act, x, b))
+
+    function ∇biasact(act, x, b)
+        dx = Enzyme.make_zero(x)
+        db = Enzyme.make_zero(b)
+        Enzyme.autodiff(
+            Reverse,
+            sumabs2biasact,
+            Active,
+            Const(act),
+            Duplicated(x, dx),
+            Duplicated(b, db),
+        )
+        return dx, db
+    end
+
+    function ∇biasact!!(act, x, b)
+        dx = Enzyme.make_zero(x)
+        db = Enzyme.make_zero(b)
+        Enzyme.autodiff(
+            Reverse,
+            sumabs2biasact!!,
+            Active,
+            Const(act),
+            Duplicated(x, dx),
+            Duplicated(b, db),
+        )
+        return dx, db
+    end
+
+    @testset for act in (identity, relu, sigmoid, tanh, gelu)
+        x = randn(Float32, 10, 10)
+        b = randn(Float32, 10)
+
+        x_ra = Reactant.ConcreteRArray(x)
+        b_ra = Reactant.ConcreteRArray(b)
+
+        f_compile = Reactant.compile(biasact, (act, x_ra, b_ra))
+        f_compile!! = Reactant.compile(biasact!!, (act, x_ra, b_ra))
+
+        y_simple = biasact(act, x, b)
+        y_simple!! = biasact!!(act, x, b)
+        y_compile = f_compile(act, x_ra, b_ra)
+        y_compile!! = f_compile!!(act, x_ra, b_ra)
+
+        @test y_simple ≈ y_compile broken = (act === gelu)
+        @test y_simple!! ≈ y_compile!! broken = (act === gelu)
+
+        # FIXME: Seems broken currently
+        @testset "Enzyme: bias_activation" begin
+            ∂x_enz, ∂b_enz = ∇biasact(act, x, b)
+            ∇biasact_compiled = Reactant.compile(∇biasact, (act, x_ra, b_ra))
+            ∂x_compile, ∂b_compile = ∇biasact_compiled(act, x_ra, b_ra)
+
+            @test ∂x_enz ≈ ∂x_compile broken = (act === gelu)
+            @test ∂b_enz ≈ ∂b_compile broken = (act === gelu)
+        end
+
+        @testset "Enzyme: bias_activation!!" begin
+            ∂x_enz!!, ∂b_enz!! = ∇biasact!!(act, x, b)
+            ∇biasact!!_compiled = Reactant.compile(∇biasact!!, (act, x_ra, b_ra))
+            ∂x_compile!!, ∂b_compile!! = ∇biasact!!_compiled(act, x_ra, b_ra)
+
+            @test ∂x_enz!! ≈ ∂x_compile!! broken = (act === gelu)
+            @test ∂b_enz!! ≈ ∂b_compile!! broken = (act === gelu)
+        end
+    end
 end
 
 @testset "Fast Activation" begin
@@ -18,7 +88,7 @@ end
 
     function ∇sumabs2!!(f, x)
         dx = Enzyme.make_zero(x)
-        Enzyme.autodiff(Reverse, sumabs2, Active, Const(f), Duplicated(x, dx))
+        Enzyme.autodiff(Reverse, sumabs2!!, Active, Const(f), Duplicated(x, dx))
         return dx
     end
 
@@ -28,8 +98,8 @@ end
     @testset "Activation: $act" for act in (
         identity, relu, sigmoid, tanh, tanh_fast, sigmoid_fast, gelu, abs2
     )
-        f_compile = Reactant.compile(sumabs2, (act, x_act))
-        f_compile!! = Reactant.compile(sumabs2!!, (act, x_act))
+        f_compile = Reactant.compile(sumabs2, (act, x_act_ca))
+        f_compile!! = Reactant.compile(sumabs2!!, (act, x_act_ca))
 
         y_simple = sumabs2(act, x_act)
         y_simple!! = sumabs2!!(act, x_act)