Update test/mooncake.jl

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

Author: kshyatt

test/mooncake.jl

@@ -55,20 +55,27 @@ function test_pullbacks_match(rng, f!, f, A, args, Δargs, alg = nothing; rdata
 
     dargs_inplace = Δargs isa Tuple ? make_mooncake_fdata.(deepcopy(Δargs)) : make_mooncake_fdata(deepcopy(Δargs))
 
-    copy_out, copy_pb!! = isnothing(alg) ? rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_copy), Mooncake.CoDual(args, dargs_copy)) : rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_copy), Mooncake.CoDual(args, dargs_copy), Mooncake.CoDual(alg, Mooncake.NoFData()))
-    copy_pb!!(rdata)
-
-    # not every f! has a handwritten rrule!!
-    inplace_sig = isnothing(alg) ? Tuple{typeof(f!), typeof(A), typeof(args)} : Tuple{typeof(f!), typeof(A), typeof(args), typeof(alg)}
-
-    has_handwritten_rule = hasmethod(Mooncake.rrule!!, inplace_sig)
-    if has_handwritten_rule
-        inplace_out, inplace_pb!! = isnothing(alg) ? Mooncake.rrule!!(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace)) : Mooncake.rrule!!(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace), Mooncake.CoDual(alg, Mooncake.NoFData()))
+    if isnothing(alg)
+        copy_out, copy_pb!! = rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_copy), Mooncake.CoDual(args, dargs_copy))
+        inplace_sig = Tuple{typeof(f!), typeof(A), typeof(args)})
+        if hasmethod(Mooncake.rrule!!, inplace_sig) # hand written rule
+            inplace_out, inplace_pb!! = Mooncake.rrule!!(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace))
+        else
+            inplace_rrule = Mooncake.build_rrule(rvs_interp, inplace_sig)
+            inplace_out, inplace_pb!! = inplace_rrule(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace))
+        end
     else
-        rvs_interp = Mooncake.get_interpreter(Mooncake.ReverseMode)
-        inplace_rrule = Mooncake.build_rrule(rvs_interp, inplace_sig)
-        inplace_out, inplace_pb!! = isnothing(alg) ? inplace_rrule(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace)) : inplace_rrule(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace), Mooncake.CoDual(alg, Mooncake.NoFData()))
+        copy_out, copy_pb!! = rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_copy), Mooncake.CoDual(args, dargs_copy), Mooncake.CoDual(alg, Mooncake.NoFData()))
+        inplace_sig = Tuple{typeof(f!), typeof(A), typeof(args), typeof(alg)}
+        if hasmethod(Mooncake.rrule!!, inplace_sig) # hand written rule
+            inplace_out, inplace_pb!! = Mooncake.rrule!!(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace), Mooncake.CoDual(alg, Mooncake.NoFData()))
+        else
+            inplace_rrule = Mooncake.build_rrule(rvs_interp, inplace_sig)
+            inplace_out, inplace_pb!! = inplace_rrule(Mooncake.CoDual(f!, Mooncake.NoFData()), Mooncake.CoDual(Ac, dA_inplace), Mooncake.CoDual(args, dargs_inplace), Mooncake.CoDual(alg, Mooncake.NoFData()))
+        end
     end
+        
+    copy_pb!!(rdata)
     inplace_pb!!(rdata)
 
     dA_inplace_ = Mooncake.arrayify(A, dA_inplace)[2]