Require backends

[nmheim]

This PR tries to separate the different AD backends into files that can be loaded optionally via Requires.jl.

TODO

• univariate distributions
• multivariate distributions
• matrix variate distributions
• product distributions
• test each backend separately
• ForwardDiff tests for test/others.jl
• move adjoints in zygote.jl back to original functions definitions

[devmotion]

Thanks for the PR, and IMO it's much easier to review properly in this way! I added some preliminary comments (I haven't checked all details yet).

[devmotion]

Is it necessary to put it in a separate module? I would be surprised at least if the use of Requires and separate modules are needed to define the adjoint for Tracker.

I would have assumed that something like

if AD == "All" || AD == "Tracker"
    include("tracker_utils.jl")
end

in test/ad/distributions.jl would be sufficient.

[nmheim]

The problem is that the TrackedMatrix overload is compiled even if Tracker is not loaded...
I tried it with Requires (and with include) in test/ad/distributions.jl which didn't work, but defining is_posdef in test/runtests.jl and requiring the Tracker stuff there does the job. So there is no more separate module, but there might be an obvious solution I am not seeing?

[nmheim]

There is a file src/for.jl that is not loaded in src/DistributionsAD.jl, that was introduced in this PR: #13

what to do with it?

[yebai]

There is a file src/for.jl that is not loaded in src/DistributionsAD.jl, that was introduced in this PR: #13
what to do with it?

@nmheim These codes should probably live in the package Probability.jl. But for now, perhaps consider moving these code into a dedicated (temporary) subfolder, say combinators/?

[sethaxen]

I had a few suggestions for the ChainRules overloads

[sethaxen]

Is there anything in ChainRules specifically that you need, or could you use ChainRulesCore, which is much lighter? There's also ChainRulesTestUtils.rrule_test, which makes it easy to test the rules directly.

[nmheim]

First, thanks a lot for all the suggestions! :)

ChainRules is needed for now, because we are using chol_blocked_rev here:

DistributionsAD.jl/src/common.jl

Line 11 in 5ec5201

∂X = @thunk(ChainRules.chol_blocked_rev(f̄, factors, 25, true))

[sethaxen]

You're welcome! That makes sense. I added an alternate suggestion below.

[sethaxen]

SInce you define turing_chol as

function turing_chol(A::AbstractMatrix, check)
    chol = cholesky(A, check=check)
    (chol.factors, chol.info)
end

couldn't you define its rrule as (something like)

function ChainRulesCore.rrule(::typeof(turing_chol), A::AbstractMatrix, check)
    chol, inner_pullback = rrule(cholesky, A, check=check)
    Ω = (chol.factors, chol.info)
    function turing_chol_pullback(ΔΩ)
        (Δfactors, Δinfo) = ΔΩ
        Δchol = Composite{typeof(chol)}(factors = Δfactors, info = Δinfo)
        (_, ∂A, ∂check) = inner_pullback(Δchol)
        return (NO_FIELDS, A, ∂check)
    end
    return Ω, turing_chol_pullback
end

? This would be a bit safer, as it doesn't use ChainRules internal functions.

The check argument isn't currently supported by ChainRules. I'll fix.

[sethaxen]

Currently chol_blocked_rev is only tested for real matrices. Probably the right thing to do is relax that constraint on ChainRules and test (I'll open a PR).

[nmheim]

I will push the suggestions by @sethaxen once ChainRules can deal with keyword arguments. Then this should be ready.

[nmheim]

Ok, it will take some more time until the keyword stuff for the cholesky rule is fixed. Seth will ping me once thats done. I think we could merge this? @devmotion @mohamed82008