Handle mixed types for inner properly

Author: kshyatt

ext/VectorInterfaceEnzymeExt.jl

@@ -215,6 +215,16 @@ function EnzymeRules.forward(
     end
 end
 
+function project_add!(C, A, α)
+    TC = Base.promote_op(+, scalartype(A), scalartype(α))
+    return if !(TC <: Real) && scalartype(C) <: Real
+        add!(C, real(add!(zerovector(C, TC), A, α)))
+    else
+        add!(C, A, α)
+    end
+end
+
+
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfigWidth{1},
         func::Const{typeof(inner)},
@@ -241,8 +251,8 @@ function EnzymeRules.reverse(
     )
     ΔS = dret.val
     Aval, Bval = cache
-    !isa(A, Const) && add!(A.dval, Bval, conj(ΔS))
-    !isa(B, Const) && add!(B.dval, Aval, ΔS)
+    !isa(A, Const) && project_add!(A.dval, Bval, conj(ΔS))
+    !isa(B, Const) && project_add!(B.dval, Aval, ΔS)
     return (nothing, nothing)
 end
 

ext/VectorInterfaceMooncakeExt.jl

@@ -140,6 +140,15 @@ end
 # inner
 # -----
 
+function project_add!(C, A, α)
+    TC = Base.promote_op(+, scalartype(A), scalartype(α))
+    return if !(TC <: Real) && scalartype(C) <: Real
+        add!(C, real(add!(zerovector(C, TC), A, α)))
+    else
+        add!(C, A, α)
+    end
+end
+
 @is_primitive DefaultCtx Tuple{typeof(inner), AbstractArray, AbstractArray}
 
 function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual, B_ΔB::CoDual)
@@ -151,8 +160,8 @@ function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual, B_ΔB::CoDual)
     s = inner(A, B)
 
     function inner_pullback(Δs)
-        add!(ΔA, B, conj(Δs))
-        add!(ΔB, A, Δs)
+        project_add!(ΔA, B, conj(Δs))
+        project_add!(ΔB, A, Δs)
         return NoRData(), NoRData(), NoRData()
     end
 

test/enzyme.jl

@@ -104,13 +104,13 @@ end
     end
 end
 
-@testset "inner ($T)" for T in eltypes
+@testset "inner ($Tx, $Ty)" for Tx in eltypes, Ty in eltypes
     n = 12
-    atol = rtol = n * precision(T)
+    atol = rtol = n * max(precision(Tx), precision(Ty))
 
     # Vector
-    x = randn(T, n)
-    y = randn(T, n)
+    x = randn(Tx, n)
+    y = randn(Ty, n)
     for RT in (Const, Active)
         test_reverse(inner, RT, (x, Duplicated), (y, Duplicated); atol, rtol)
     end

test/mooncake.jl

@@ -73,13 +73,13 @@ end
     Mooncake.TestUtils.test_rule(rng, add!!, my, mx, α, β; atol, rtol, is_primitive = false)
 end
 
-@testset "inner pullbacks ($T)" for T in eltypes
+@testset "inner pullbacks ($Tx, $Ty)" for Tx in eltypes, Ty in eltypes
     n = 12
-    atol = rtol = n * precision(T)
+    atol = rtol = n * max(precision(Tx), precision(Ty))
 
     # Vector
-    x = randn(T, n)
-    y = randn(T, n)
+    x = randn(Tx, n)
+    y = randn(Ty, n)
     Mooncake.TestUtils.test_rule(rng, inner, x, y; atol, rtol, is_primitive = false)
 
     # MinimalMVec