attempt to fix derivative tests

Author: lxvm

ext/IntegralsZygoteExt.jl

@@ -14,11 +14,19 @@ ChainRulesCore.@non_differentiable Integrals.checkkwargs(kwargs...)
 ChainRulesCore.@non_differentiable Integrals.isinplace(f, args...)    # fixes #99
 ChainRulesCore.@non_differentiable Integrals.init_cacheval(alg, prob)
 
-function ChainRulesCore.rrule(::typeof(Integrals.transformation_if_inf), f, domain)
-    function transformation_if_inf_pullback(Δ)
-        return NoTangent(), Δ...
+function ChainRulesCore.rrule(::typeof(Integrals.__solve), cache::Integrals.IntegralCache,
+        alg::Integrals.ChangeOfVariables, sensealg, udomain, p;
+        kwargs...)
+    _cache, vdomain = Integrals._change_variables(cache, alg, sensealg, udomain, p)
+    sol, back = Zygote.pullback((args...) -> Integrals.__solve(args...; kwargs...),
+        _cache, alg.alg, sensealg, vdomain, p)
+    function change_of_variables_pullback(Δ)
+        return (NoTangent(), back(Δ)...)
     end
-    return Integrals.transformation_if_inf(f, domain), transformation_if_inf_pullback
+    prob = Integrals.build_problem(cache)
+    _sol = SciMLBase.build_solution(
+        prob, alg.alg, sol.u, sol.resid, chi = sol.chi, retcode = sol.retcode, stats = sol.stats)
+    return _sol, change_of_variables_pullback
 end
 
 # we will need to implement the following adjoints when we compute ∂f/∂u

src/Integrals.jl

@@ -78,6 +78,14 @@ end
 
 function __solve(cache::IntegralCache, alg::ChangeOfVariables, sensealg, udomain, p;
         kwargs...)
+    _cache, vdomain = _change_variables(cache, alg, sensealg, udomain, p)
+    sol = __solve(_cache, alg.alg, sensealg, vdomain, p; kwargs...)
+    prob = build_problem(cache)
+    return SciMLBase.build_solution(
+        prob, alg.alg, sol.u, sol.resid, chi = sol.chi, retcode = sol.retcode, stats = sol.stats)
+end
+
+function _change_variables(cache, alg, sensealg, udomain, p)
     cacheval = cache.cacheval.alg
     g, vdomain = alg.fu2gv(cache.f, udomain)
     _cache = IntegralCache(Val(isinplace(g)),
@@ -89,7 +97,7 @@ function __solve(cache::IntegralCache, alg::ChangeOfVariables, sensealg, udomain
         sensealg,
         cache.kwargs,
         cacheval)
-    return __solve(_cache, alg.alg, sensealg, vdomain, p; kwargs...)
+    return _cache, vdomain
 end
 
 function get_prototype(prob::IntegralProblem)