fixup

ext/QuadGKEnzymeExt.jl

@@ -67,17 +67,17 @@ end
     return rt == Enzyme.Compiler.AnyState || rt == Enzyme.Compiler.DupState
 end
 
-function Base.:+(a::ClosureVector, b::ClosureVector)
-    Enzyme.Compiler.recursive_add(a, b, identity, guaranteed_nonactive)
+function Base.:+(a::CV, b::CV) where {CV <: ClosureVector}
+    Enzyme.Compiler.recursive_add(a, b, identity, guaranteed_nonactive)::CV
 end
 
-function Base.:-(a::ClosureVector, b::ClosureVector)
-    Enzyme.Compiler.recursive_add(a, b, x->-x, guaranteed_nonactive)
+function Base.:-(a::CV, b::CV) where {CV <: ClosureVector}
+    Enzyme.Compiler.recursive_add(a, b, x->-x, guaranteed_nonactive)::CV
 end
 
-function Base.:*(a::Number, b::ClosureVector)
+function Base.:*(a::Number, b::CV) where {CV <: ClosureVector}
     # b + (a-1) * b = a * b
-    Enzyme.Compiler.recursive_add(b, b, x->(a-1)*x, guaranteed_nonactive)
+    Enzyme.Compiler.recursive_add(b, b, x->(a-1)*x, guaranteed_nonactive)::CV
 end
 
 function Base.:*(a::ClosureVector, b::Number)
@@ -105,4 +105,32 @@ function Enzyme.EnzymeRules.reverse(config, ofunc::Const{typeof(quadgk)}, dres::
             dsegsn)
 end
 
+function Enzyme.EnzymeRules.reverse(config, ofunc::Const{typeof(quadgk)}, dres::Type{<:Union{Duplicated, BatchDuplicated}}, cache, f, segs::Annotation{T}...; kws...) where {T<:Real}
+    dres = cache[2]
+    @show dres
+    df = if f isa Const
+        nothing
+    else
+        segbuf = cache[1]
+        fwd, rev = Enzyme.autodiff_thunk(ReverseSplitNoPrimal, Const{typeof(call)}, Active, typeof(f), Const{T})
+        _df, _ = quadgk(map(x->x.val, segs)...; kws..., eval_segbuf=segbuf, maxevals=0, norm=f->0) do x
+            @show x
+            tape, prim, shad = fwd(Const(call), f, Const(x))
+            @show prim, shad
+            shad .= dres
+            drev = rev(Const(call), f, Const(x), tape)
+            @show drev
+            return ClosureVector(drev[1][1])
+        end
+        _df.f
+    end
+    dsegs1 = segs[1] isa Const ? nothing : -LinearAlgebra.dot(f.val(segs[1].val), dres)
+    dsegsn = segs[end] isa Const ? nothing : LinearAlgebra.dot(f.val(segs[end].val), dres)
+    Enzyme.make_zero!(dres)
+    return (df, # f
+            dsegs1,
+            ntuple(i -> nothing, Val(length(segs)-2))...,
+            dsegsn)
+end
+
 end # module

test/runtests.jl

@@ -434,9 +434,26 @@ end
     f2(x) = quadgk(cos, x, 1)[1]
     f3(x) = quadgk(y->cos(x * y), 0., 1.)[1]
 
+    f1_count(x) = quadgk_count(cos, 0., x)[1]
+    f2_count(x) = quadgk_count(cos, x, 1)[1]
+    f3_count(x) = quadgk_count(y->cos(x * y), 0., 1.)[1]
+
+    f_vec(x) = sum(quadgk(y->[cos(x[1] * y), cos(x[2] * y)], 0., 1.)[1])
+
     @testset "Enzyme" begin
         @test cos(0.3) ≈ Enzyme.autodiff(Reverse, f1, Active(0.3))[1][1]
         @test -cos(0.3) ≈ Enzyme.autodiff(Reverse, f2, Active(0.3))[1][1]
         @test (0.3 * cos(0.3) - sin(0.3))/(0.3*0.3) ≈ Enzyme.autodiff(Reverse, f3, Active(0.3))[1][1]
+
+        @test cos(0.3) ≈ Enzyme.autodiff(Reverse, f1_count, Active(0.3))[1][1]
+        @test -cos(0.3) ≈ Enzyme.autodiff(Reverse, f2_count, Active(0.3))[1][1]
+        @test_broken (0.3 * cos(0.3) - sin(0.3))/(0.3*0.3) ≈ Enzyme.autodiff(Reverse, f3_count, Active(0.3))[1][1]
+
+        x = [0.3, 0.7]
+        dx = [0.0, 0.0]
+        f_vec(x)
+        # TODO custom rule with mixed vector returns not yet supported
+        @test_throws Enzyme.Compiler.EnzymeRuntimeException autodiff(Reverse, f_vec, Duplicated(x, dx))
+        # @test dx ≈ [(0.3 * cos(0.3) - sin(0.3))/(0.3*0.3), (0.7 * cos(0.7) - sin(0.7))/(0.7*0.7)]
     end
 end