fix

Author: wsmoses

lib/EnzymeTestUtils/src/EnzymeTestUtils.jl

@@ -15,6 +15,7 @@ include("test_approx.jl")
 include("compatible_activities.jl")
 include("finite_difference_calls.jl")
 include("generate_tangent.jl")
+include("test_utils.jl")
 include("test_forward.jl")
 include("test_reverse.jl")
 

lib/EnzymeTestUtils/src/test_forward.jl

@@ -63,15 +63,15 @@ function test_forward(
     testset_name=nothing,
     runtime_activity::Bool=false
 )
-    call_with_copy(f, xs...) = deepcopy(f)(deepcopy(xs)...; deepcopy(fkwargs)...)
-    call_with_kwargs(f, xs...) = f(xs...; fkwargs...)
+    call_with_copy = CallWithCopyKWargs(fkwargs)
+    call_with_kwargs = CallWithKWargs(fkwargs)
     if testset_name === nothing
         testset_name = "test_forward: $f with return activity $ret_activity on $(_string_activity(args))"
     end
     @testset "$testset_name" begin
         # format arguments for autodiff and FiniteDifferences
         activities = map(Base.Fix1(auto_activity, rng), (f, args...))
-        primals = map(x -> x.val, activities)
+        primals = map(get_primal, activities)
         # call primal, avoid mutating original arguments
         fcopy = deepcopy(first(primals))
         args_copy = deepcopy(Base.tail(primals))
@@ -149,7 +149,7 @@ function test_forward(
             end
         else
             test_approx(
-                dy_ad, dy_fdm, "derivative should agree with finite differences"; atol, rtol
+                dy_ad, dy_fdm, "derivative should agree with finite differences ($activities)"; atol, rtol
             )
         end
     end

lib/EnzymeTestUtils/src/test_reverse.jl

@@ -58,8 +58,6 @@ end
 
 Here we test a rule for a function of an array in batch reverse-mode:
 
-
-
 ```julia
 x = randn(3)
 for Tret in (Const, Active), Tx in (Const, BatchDuplicated)
@@ -79,14 +77,14 @@ function test_reverse(
     testset_name=nothing,
     runtime_activity::Bool=false,
     output_tangent=nothing)
-    call_with_captured_kwargs(f, xs...) = f(xs...; fkwargs...)
+    call_with_captured_kwargs = CallWithKWargs(fkwargs)
     if testset_name === nothing
         testset_name = "test_reverse: $f with return activity $ret_activity on $(_string_activity(args))"
     end
     @testset "$testset_name" begin
         # format arguments for autodiff and FiniteDifferences
         activities = map(Base.Fix1(auto_activity, rng), (f, args...))
-        primals = map(x -> x.val, activities)
+        primals = map(get_primal, activities)
         # call primal, avoid mutating original arguments
         fcopy = deepcopy(first(primals))
         args_copy = deepcopy(Base.tail(primals))

lib/EnzymeTestUtils/src/test_utils.jl

@@ -0,0 +1,20 @@
+
+struct CallWithKWargs{KW}
+    kwargs::KW
+end
+
+function (c::CallWithKWargs)(f, xs...)
+    f(xs...; c.kwargs...)
+end
+
+struct CallWithCopyKWargs{KW}
+    kwargs::KW
+end
+
+function (c::CallWithCopyKWargs)(f, xs...)
+    deepcopy(f)(deepcopy(xs)...; deepcopy(c.kwargs)...)
+end
+
+@inline function get_primal(x::Annotation)
+    x.val
+end

src/rules/customrules.jl

@@ -48,8 +48,10 @@ import LinearAlgebra
             :(LinearAlgebra.dot(partial,dx))
         elseif dx <: AbstractFloat || dx <: AbstractArray{<:AbstractFloat}
             :(LinearAlgebra.dot(dx, partial))
-        else
+        elseif partial <: AbstractVector
             :(LinearAlgebra.dot(adjoint(partial),dx))
+        else
+            :(LinearAlgebra.dot(conj(partial),dx))
         end
         return quote
             Base.@_inline_meta
@@ -106,10 +108,52 @@ import LinearAlgebra
     end
 end
 
-@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial, dx)
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::Real, dx)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, partial, dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::Complex, dx)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, conj(partial), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Real}, dx::Number)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, partial, dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Complex}, dx::Number)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, conj(partial), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Real, N}, dx::AbstractArray{<:Any, N}) where N
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, partial, dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Complex, N}, dx::AbstractArray{<:Any, N}) where N
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, conj(partial), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractVector{<:Complex}, dx::AbstractVector{<:Any})
     EnzymeCore.EnzymeRules.multiply_fwd_into(prev, adjoint(partial), dx)
 end
 
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractMatrix{<:Real}, dx::AbstractVector)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, transpose(partial), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractMatrix{<:Complex}, dx::AbstractVector)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, adjoint(partial), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Real}, dx::AbstractArray)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, Base.permutedims(partial, (((ndims(dx)+1):ndims(partial))..., Base.OneTo(ndims(dx))...)), dx)
+end
+
+@inline function EnzymeCore.EnzymeRules.multiply_rev_into(prev, partial::AbstractArray{<:Complex}, dx::AbstractArray)
+    pd = Base.permutedims(partial, (((ndims(dx)+1):ndims(partial))..., Base.OneTo(ndims(dx))...))
+    Base.conj!(pd)
+    EnzymeCore.EnzymeRules.multiply_fwd_into(prev, pd, dx)
+end
+
 function enzyme_custom_setup_args(
     @nospecialize(B::Union{Nothing, LLVM.IRBuilder}),
     orig::LLVM.CallInst,

test/rules/internal_rules.jl

@@ -220,7 +220,13 @@ end
         end
     end
     @testset "reverse" begin
-        @testset for RT in (Const, Active,), Tx in (Const, Duplicated,),
+        @testset for RT in (Const, Active,), Tx in (Const, Duplicated,)
+
+            # TODO see https://github.com/EnzymeAD/Enzyme/issues/2537
+            if RT <: Const
+                continue
+            end
+
             x = [4.0 3.0; 2.0 1.0]
             test_reverse(LinearAlgebra.det, RT, (x, Tx))