Autodiff of `quadgk`

[Vilin97]

Thank you for a great package!
I would like to use Zygote on quadgk but below is a self-contained example, where it does not work because quadgk calls setindex!. Can I do anything to work around this? I coded up a small example of differentiating quadgk of a constant function and it did work so it seems that there is hope for Zygote and QuadGK to play together nicely.

using Zygote, QuadGK

"P(0<X<x)"
function F(c, x)
    if x < c
        return 0.
    elseif x < 1.
        return exp(c)*(x-c)^2/2
    elseif x < 1. + c
        return exp(c)*(1-c)*(2x - c - 1.)/2.
    elseif x < 2.
        return exp(c)*(4x - x^2 - 2c - 2)/2.
    else
        return exp(c)*(1-c)
    end
end

"pdf of X"
function f(c, x)
    if x < c || x > 2.
        return 0.
    elseif x < 1.
        return exp(c)*(x-c)
    elseif x < 1. + c
        return exp(c)*(1-c)
    else
        return exp(c)*(2-x)
    end
end

function D(c)
    exp(c)*(1-c)
end

function W(c1, c2)
    val, acc = quadgk(t -> F(c1, t)*f(c2, t), c1, 2)
    val
end

function prob_of_win(c1, c2)
    ((1-D(c1))*D(c2) + W(c2, c1))/(D(c1) + D(c2) - D(c1)*D(c2))
end

gradient(prob_of_win, 0.25, 0.25) # ERROR: Mutating arrays is not supported -- called setindex!(Vector{QuadGK.Segment{Float64, Float64, Float64}}, ...)

[stevengj]

Basically this needs to be done by writing ChainRules that tell Zygote how to differentiate integrals without differentiating through the (mutating) code.

Fortunately, this has already been done — if you use https://github.com/SciML/Integrals.jl, then it provides a wrapper around QuadGK that provides the appropriate ChainRules.

In the longer run it might be worth adding chain rules to QuadGK directly.

[lxvm]

Autodiff in Integrals.jl was recently fixed in this PR SciML/Integrals.jl#175 and it should be available in its next release.

[stevengj]

Note that we have manual Enzyme rules, but it would be nice to add ChainRules too.