Merge pull request #319 from AstitvaAggarwal/dev

Fix CI AD tests

Author: ChrisRackauckas

Project.toml

@@ -61,7 +61,7 @@ HCubature = "1.7"
 LinearAlgebra = "1.10"
 MCIntegration = "0.4.2"
 MonteCarloIntegration = "0.2"
-Mooncake = "0.4.184, 0.5"
+Mooncake = "0.5.6"
 Pkg = "1.10"
 QuadGK = "2.11"
 Random = "1.10"

ext/IntegralsMooncakeExt.jl

@@ -2,9 +2,8 @@ module IntegralsMooncakeExt
 using Mooncake
 using LinearAlgebra: dot
 using Integrals, SciMLBase, QuadGK
-using Mooncake: @from_chainrules, @is_primitive, increment!!, MinimalCtx, rrule!!, NoFData,
-    NoRData, CoDual, primal, NoRData, zero_fcodual
-import Mooncake: increment_and_get_rdata!, @zero_derivative
+using Mooncake: @from_chainrules, @is_primitive, increment!!, MinimalCtx, NoFData,
+    CoDual, primal, NoRData, zero_fcodual, increment_and_get_rdata!, @zero_derivative
 using Integrals: AbstractIntegralMetaAlgorithm, IntegralProblem
 import ChainRulesCore
 import ChainRulesCore: Tangent, NoTangent, ProjectTo
@@ -34,42 +33,9 @@ batch_unwrap(x::AbstractArray) = dropdims(x; dims = ndims(x))
     Union{<:AbstractVector, <:Number}, Union{<:AbstractVector, <:Number},
 }
 
-# @from_chainrules MinimalCtx Tuple{Type{IntegralProblem{iip}},Any,Any,Any} where {iip} true
-@is_primitive MinimalCtx Tuple{Type{IntegralProblem{iip}}, Any, Any, Any} where {iip}
-function Mooncake.rrule!!(
-        ::CoDual{Type{IntegralProblem{iip}}}, f::CoDual,
-        domain::CoDual, p::CoDual; kwargs...
-    ) where {iip}
-    f_prim, domain_prim, p_prim = map(primal, (f, domain, p))
-    prob = IntegralProblem{iip}(f_prim, domain_prim, p_prim; kwargs...)
-
-    function IntegralProblem_iip_pullback(Δ)
-        data = Δ isa NoRData ? Δ : Δ.data
-        ddomain = hasproperty(data, :domain) ? data.domain : NoRData()
-        dp = hasproperty(data, :p) ? data.p : NoRData()
-        dkwargs = hasproperty(Δ, :kwargs) ? data.kwargs : NoRData()
-
-        # domain is always a Tuple, so it always has NoFData
-        # below conditional is in case p is an Array or similar
-        if Mooncake.rdata_type(typeof(p_prim)) == NoRData()
-            Mooncake.increment!!(p.dx, dp)
-            grad_p = NoRData()
-        else
-            grad_p = dp
-        end
-
-        return NoRData(), NoRData(), ddomain, grad_p, dkwargs
-    end
-    return zero_fcodual(prob), IntegralProblem_iip_pullback
-end
-
-# Mooncake does not need chainrule for evaluate! as it supports mutation.
+@from_chainrules MinimalCtx Tuple{Type{IntegralProblem{iip}}, Any, Any, Any} where {iip} true
 @from_chainrules MinimalCtx Tuple{Type{IntegralProblem}, Any, Any, Any} true
 @from_chainrules MinimalCtx Tuple{typeof(Integrals.u2t), Any, Any} true
-@from_chainrules MinimalCtx Tuple{
-    typeof(SciMLBase.build_solution), IntegralProblem, Any, Any, Any,
-} true
-
 @from_chainrules MinimalCtx Tuple{typeof(Integrals.__solvebp), Any, Any, Any, Any, Any} true
 
 # Add MooncakeVJP support to the dispatch function defined in ZygoteExt
@@ -112,135 +78,4 @@ function Integrals._compute_dfdp_and_f(::Integrals.MooncakeVJP, cache, p, Δ)
     return dfdp, _f
 end
 
-# Internal Mooncake overloads to accommodate IntegralSolution etc. Struct's Tangent Types.
-# Allows clear translation from ChainRules -> Mooncake's tangent.
-function Mooncake.increment_and_get_rdata!(
-        f::NoFData, r::Tuple{T, T},
-        t::Union{Tangent{Tuple{T, T}, Tuple{T, T}}, Tangent{Any, Tuple{Float64, Float64}}}
-    ) where {T <: Base.IEEEFloat}
-    return r .+ t.backing
-end
-
-function Mooncake.increment_and_get_rdata!(
-        f::Tuple{Vector{T}, Vector{T}},
-        r::NoRData,
-        t::Tangent{Any, Tuple{Vector{T}, Vector{T}}}
-    ) where {T <: Base.IEEEFloat}
-    Mooncake.increment!!(f, t.backing)
-    return NoRData()
-end
-
-# sol.u & p are single scalar values, domain (lb,ub) is single/multi - variate.
-function Mooncake.increment_and_get_rdata!(
-        f::NoFData,
-        r::T,
-        t::Tangent{
-            Any,
-            @NamedTuple{
-                u::T,
-                resid::R,
-                prob::Tangent{
-                    Any,
-                    @NamedTuple{
-                        f::NoTangent,
-                        domain::Tangent{Any, Tuple{M, M}},
-                        p::P,
-                        kwargs::NoTangent,
-                    }
-                },
-                alg::A,
-                retcode::NoTangent,
-                chi::NoTangent,
-                stats::NoTangent,
-            }
-        }
-    ) where {
-        T <: Base.IEEEFloat,
-        R <: Union{NoTangent, T},
-        P <: Union{T, Vector{T}},
-        M <: Union{T, Vector{T}},
-        A <: Union{
-            NoTangent,
-            Tangent{
-                Any,
-                @NamedTuple{
-                    nodes::Vector{T},
-                    weights::Vector{T},
-                    subintervals::NoTangent,
-                }
-            }
-        }
-    }
-    # rdata component of t + r (u field)
-    return Mooncake.increment_and_get_rdata!(f, r, t.u)
-end
-
-# sol.u is vector valued, p is scalar/vector valued, domain can be single/multi - variate
-#  resid can be single/vector valued. For inplace integrals (iip true) : included integrand_prototype field in typeof{prob.f}
-function Mooncake.increment_and_get_rdata!(
-        f::Vector{T},
-        r::NoRData,
-        t::Union{
-            Tangent{
-                Any,
-                @NamedTuple{
-                    u::Vector{T},
-                    resid::R,
-                    prob::Tangent{
-                        Any,
-                        @NamedTuple{
-                            f::F,
-                            domain::Tangent{Any, M},
-                            p::P,
-                            kwargs::NoTangent,
-                        }
-                    },
-                    alg::A,
-                    retcode::NoTangent,
-                    chi::NoTangent,
-                    stats::NoTangent,
-                }
-            }
-        }
-    ) where {
-        T <: Base.IEEEFloat,
-        R <: Union{NoTangent, T, Vector{T}},
-        P <: Union{T, Vector{T}},
-        M <: Union{Tuple{T, T}, Tuple{Vector{T}, Vector{T}}},
-        F <: Union{
-            NoTangent,
-            Tangent{
-                Any,
-                @NamedTuple{
-                    f::NoTangent,
-                    integrand_prototype::Vector{T},
-                }
-            }
-        },
-        A <: Union{
-            NoTangent,
-            Tangent{
-                Any,
-                @NamedTuple{
-                    nodes::Vector{T},
-                    weights::Vector{T},
-                    subintervals::NoTangent,
-                }
-            }
-        }
-    }
-    Mooncake.increment!!(f, t.u)
-    # rdata component(t) + r
-    return t.prob.domain
-end
-
-# cannot mutate NoRData() in place, therefore return as is.
-function Mooncake.increment!!(
-        ::Mooncake.NoRData,
-        y::Tangent{Any, Y}
-    ) where {
-        T <: Base.IEEEFloat, Y <: Union{Tuple{T, T}, Tuple{Vector{T}, Vector{T}}},
-    }
-    return Mooncake.NoRData()
-end
 end

test/derivative_tests.jl

@@ -182,7 +182,6 @@ do_tests_mooncake = function (; f, scalarize, lb, ub, p, alg, abstol, reltol)
     end
     sol_fp = testf(lb, ub, p)
 
-    # sensealg when non zygoet?
     cache = Mooncake.prepare_gradient_cache(
         testf, lb, ub, p isa Number && f isa BatchIntegralFunction ? Scalar(p) : p
     )