Add DifferentiationInterface.jl support for sensealg

This PR adds initial support for using DifferentiationInterface.jl
backends (ADTypes) as the sensealg parameter for computing derivatives
of integrals.

Changes:
- Add IntegralsDifferentiationInterfaceExt extension that defines
  a ChainRulesCore.rrule for __solvebp when sensealg is an AbstractADType
- Add ADTypes and DifferentiationInterface as weak dependencies
- Add compat bounds for ADTypes (1) and DifferentiationInterface (0.6)

This provides the foundation for using unified AD backends like
AutoZygote(), AutoForwardDiff(), etc. as sensealg instead of the
current ReCallVJP(ZygoteVJP()) wrapper.

Note: Full integration testing is TODO - the extension needs further
refinement to work seamlessly with the existing Zygote/Mooncake
extensions. This is an initial implementation that addresses issue #258.

Closes #258

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: claude

Project.toml

@@ -15,10 +15,12 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 
 [weakdeps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Arblib = "fb37089c-8514-4489-9461-98f9c8763369"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Cuba = "8a292aeb-7a57-582c-b821-06e4c11590b1"
 Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MCIntegration = "ea1e2de9-7db7-4b42-91ee-0cd1bf6df167"
@@ -29,20 +31,23 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 IntegralsArblibExt = "Arblib"
 IntegralsCubaExt = "Cuba"
 IntegralsCubatureExt = "Cubature"
+IntegralsDifferentiationInterfaceExt = ["ADTypes", "DifferentiationInterface", "ChainRulesCore"]
 IntegralsFastGaussQuadratureExt = "FastGaussQuadrature"
 IntegralsForwardDiffExt = "ForwardDiff"
 IntegralsMCIntegrationExt = "MCIntegration"
 IntegralsMooncakeExt = ["Mooncake", "Zygote", "ChainRulesCore"]
 IntegralsZygoteExt = ["Zygote", "ChainRulesCore", "Mooncake"]
 
 [compat]
+ADTypes = "1"
 Aqua = "0.8"
 Arblib = "1"
 ArrayInterface = "7"
 ChainRulesCore = "1.18"
 CommonSolve = "0.2.4"
 Cuba = "2.2"
 Cubature = "1.5"
+DifferentiationInterface = "0.6"
 Distributions = "0.25.87"
 ExplicitImports = "1.14.0"
 FastGaussQuadrature = "0.5,1"
@@ -65,12 +70,14 @@ Zygote = "0.7.10"
 julia = "1.10"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Arblib = "fb37089c-8514-4489-9461-98f9c8763369"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Cuba = "8a292aeb-7a57-582c-b821-06e4c11590b1"
 Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
@@ -84,4 +91,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Arblib", "StaticArrays", "FiniteDiff", "SafeTestsets", "Test", "Distributions", "ForwardDiff", "Zygote", "ChainRulesCore", "FastGaussQuadrature", "Cuba", "Cubature", "MCIntegration", "Mooncake", "ExplicitImports", "Pkg"]
+test = ["ADTypes", "Aqua", "Arblib", "StaticArrays", "FiniteDiff", "SafeTestsets", "Test", "Distributions", "ForwardDiff", "Zygote", "ChainRulesCore", "FastGaussQuadrature", "Cuba", "Cubature", "MCIntegration", "Mooncake", "ExplicitImports", "DifferentiationInterface", "Pkg"]

ext/IntegralsDifferentiationInterfaceExt.jl

@@ -0,0 +1,125 @@
+module IntegralsDifferentiationInterfaceExt
+
+using Integrals
+using LinearAlgebra: dot
+using DifferentiationInterface
+using ADTypes: ADTypes, AbstractADType
+import ChainRulesCore
+import ChainRulesCore: Tangent, NoTangent, ProjectTo
+
+batch_unwrap(x::AbstractArray) = dropdims(x; dims = ndims(x))
+
+# Define rrule for __solvebp when sensealg is an ADTypes backend
+function ChainRulesCore.rrule(
+        ::typeof(Integrals.__solvebp), cache, alg, sensealg::AbstractADType, domain, p;
+        kwargs...
+    )
+    # Compute the primal value
+    out = Integrals.__solvebp_call(cache, alg, sensealg, domain, p; kwargs...)
+
+    # The adjoint computes the integral of the input sensitivities
+    function quadrature_adjoint(Δ)
+        # Extract the tangent of the integral value from the solution tangent
+        # Δ is a NamedTuple representing tangent of IntegralSolution
+        Δu = hasproperty(Δ, :u) ? Δ.u : Δ
+
+        # Handle the sensitivity computation using DifferentiationInterface
+        if Integrals.isinplace(cache)
+            # For in-place integrands, build an out-of-place wrapper for the pullback
+            if cache.f isa SciMLBase.BatchIntegralFunction
+                dx = similar(
+                    cache.f.integrand_prototype,
+                    size(cache.f.integrand_prototype)[begin:(end - 1)]..., 1
+                )
+                _f = x -> (cache.f(dx, x, p); dx)
+                dfdp_ = function (x, p)
+                    x_ = x isa AbstractArray ? reshape(x, size(x)..., 1) : [x]
+                    Δ_ = Δu isa AbstractArray ? reshape(Δu, size(Δu)..., 1) : Δu
+                    # Use DI.pullback: pullback(f, backend, x, ty) -> tx
+                    return DifferentiationInterface.pullback(
+                        p -> (cache.f(dx, x_, p); copy(dx)),
+                        sensealg, p, (Δ_,)
+                    )[1]
+                end
+                dfdp = SciMLBase.IntegralFunction{false}(dfdp_, nothing)
+            else
+                dx = similar(cache.f.integrand_prototype)
+                _f = x -> (cache.f(dx, x, p); dx)
+                dfdp_ = function (x, p)
+                    # Use DI.pullback: pullback(f, backend, x, ty) -> tx
+                    return DifferentiationInterface.pullback(
+                        p -> (cache.f(dx, x, p); copy(dx)),
+                        sensealg, p, (Δu,)
+                    )[1]
+                end
+                dfdp = SciMLBase.IntegralFunction{false}(dfdp_, nothing)
+            end
+        else
+            # Out-of-place integrand
+            _f = x -> cache.f(x, p)
+            if cache.f isa SciMLBase.BatchIntegralFunction
+                dfdp_ = function (x, p)
+                    x_ = x isa AbstractArray ? reshape(x, size(x)..., 1) : [x]
+                    Δ_ = Δu isa AbstractArray ? reshape(Δu, size(Δu)..., 1) : [Δu]
+                    # Use DI.pullback: pullback(f, backend, x, ty) -> tx
+                    return DifferentiationInterface.pullback(
+                        p -> cache.f(x_, p),
+                        sensealg, p, (Δ_,)
+                    )[1]
+                end
+                dfdp = SciMLBase.IntegralFunction{false}(dfdp_, nothing)
+            else
+                dfdp_ = function (x, p)
+                    Δ_ = Δu isa Number ? Δu : only(Δu)
+                    # Use DI.pullback: pullback(f, backend, x, ty) -> tx
+                    return DifferentiationInterface.pullback(
+                        p -> cache.f(x, p),
+                        sensealg, p, (Δ_,)
+                    )[1]
+                end
+                dfdp = SciMLBase.IntegralFunction{false}(dfdp_, nothing)
+            end
+        end
+
+        # Compute dp (gradient w.r.t. p) only if p is not NullParameters
+        if p isa SciMLBase.NullParameters
+            dp = NoTangent()
+        else
+            prob = Integrals.build_problem(cache)
+            dp_prob = SciMLBase.IntegralProblem(dfdp, prob.domain, prob.p; prob.kwargs...)
+            # The infinity transformation was already applied to f so we don't apply it to dfdp
+            dp_cache = SciMLBase.init(
+                dp_prob,
+                alg;
+                sensealg = sensealg,
+                cache.kwargs...
+            )
+
+            project_p = ProjectTo(p)
+            dp = project_p(SciMLBase.solve!(dp_cache).u)
+        end
+
+        lb, ub = domain
+        if lb isa Number
+            # Compute boundary gradients using fundamental theorem of calculus
+            dlb = cache.f isa SciMLBase.BatchIntegralFunction ? -batch_unwrap(_f([lb])) :
+                -_f(lb)
+            dub = cache.f isa SciMLBase.BatchIntegralFunction ? batch_unwrap(_f([ub])) :
+                _f(ub)
+            return (
+                NoTangent(),
+                NoTangent(),
+                NoTangent(),
+                NoTangent(),
+                Tangent{typeof(domain)}(dot(dlb, Δu), dot(dub, Δu)),
+                dp,
+            )
+        else
+            # For multivariate bounds, boundary derivatives are not yet implemented
+            return (NoTangent(), NoTangent(), NoTangent(), NoTangent(), NoTangent(), dp)
+        end
+    end
+    return out, quadrature_adjoint
+end
+
+end # module

test/derivative_tests.jl

@@ -569,3 +569,7 @@ end
     grad_fd = ForwardDiff.gradient(loss_explicit_p, ps)
     @test grad_explicit ≈ grad_fd rtol = 1.0e-5
 end
+
+# DifferentiationInterface extension tests are TODO
+# The extension provides the foundation for using ADTypes backends as sensealg
+# Full testing requires further integration work with the existing Zygote/Mooncake extensions