Add Distributions integration tests

[mhauru]

Runs logpdf on a lot of different distributions, compares against FiniteDifferences.

Currently generating a lot of failures. Need to investigate which ones should be reported as bugs, which ones should just be marked as broken because of e.g. missing foreign calls.

Largely copied from similar tests in Tapir: https://github.com/compintell/Tapir.jl/blob/main/test/integration_testing/distributions.jl

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.26%. Comparing base (037dfed) to head (8071208).
⚠️ Report is 496 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1819      +/-   ##
==========================================
+ Coverage   67.50%   74.26%   +6.75%     
==========================================
  Files          31       47      +16     
  Lines       12668    16966    +4298     
==========================================
+ Hits         8552    12600    +4048     
- Misses       4116     4366     +250     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[wsmoses]

@mhauru @yebai from the quadgk version incompatability, I think this should just be done in DistributionsAD.jl?

In any case, open any issues that persist and we'll get them fixed quickly

[mhauru]

This is ready for review. I've marked all failing tests as broken and opened issues for them. The univariate distributions are doing well, except for some missing foreign calls that were already known. There's still quite a lot of breakage for multivariate and matrix variate distributions.

Here are all the relevant issues:

• Missing Rmath derivatives #1620
• Implement besseli #1603
• [ForwardMode] unimplemented derivative for trtrs #1991
• Forward mode gradient error with argument of limited mutability #1821
• Missing rule for LAPACK's dgetrf_64_ #1820
• [ForwardMode] Runtime Activity not supported for BLAS calls #1995
• LLVM error: function failed verification (4) on truncated Beta distribution #1996
• Incorrect forward mode gradient for truncated normal distribution #1998
• BatchDuplicated fails with RefValues #1997
• Forward mode gradient doesn't support tuples with varying element types #1994

EDIT:
Also

• julia crash, Illegal Instruction 4, Wishart distribution #2012
• no method matching reverse with truncated Beta distribution #2010

[mhauru]

Fixed the two issues you pointed out. Tests now locally crash julia with Illegal instruction, not sure why, probably something that has changed on main. Will investigate tomorrow.

[mhauru]

Commented out the one that crashes and made an issue about it. Tests pass locally again, ready for another look, @wsmoses.

[mhauru]

I'm unable to reproduce locally that last test failure, despite using the same Julia and dependency versions. Could something here be platform dependent?

[mhauru]

Tests pass locally for me, but hang for 6 hours on CI and time out. @wsmoses any ideas for what might be going on? Would you have a Linux machine to test this locally?

[mhauru]

@wsmoses could we see if we can make progress here? Testing Turing.jl is much more fruitful if we have confidence in individual distributions being fine. This is still in a state where tests pass for me locally but hang on CI. Could you try running them on any non-Mac machines you have available to see what happens?

[wsmoses]

honestly no clue/fruit from my end either, probably the way to debug this is to binary search in CI selectively removing half the tests each time until we find what is the one going awry

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/test/integration/Distributions/runtests.jl b/test/integration/Distributions/runtests.jl
index 1ac4435..4846671 100644
--- a/test/integration/Distributions/runtests.jl
+++ b/test/integration/Distributions/runtests.jl
@@ -45,29 +45,29 @@ end
 
 # Turn a distribution into a call to logpdf.
 function TestCase(d::Distribution, value, name, runtime_activity, broken, splat)
-    TestCase(x -> logpdf(d, x), value, name, runtime_activity, broken, splat)
+    return TestCase(x -> logpdf(d, x), value, name, runtime_activity, broken, splat)
 end
 
 # Defaults for name, runtime_activity and broken.
 function TestCase(
-    f, value;
-    name=nothing, runtime_activity=Neither, broken=Neither, splat=false
-)
+        f, value;
+        name = nothing, runtime_activity = Neither, broken = Neither, splat = false
+    )
     return TestCase(f, value, name, runtime_activity, broken, splat)
 end
 
 # Default name for a Distribution.
 function TestCase(
-    d::Distribution, value;
-    name=string(nameof(typeof(d))), runtime_activity=Neither, broken=Neither, splat=false
-)
+        d::Distribution, value;
+        name = string(nameof(typeof(d))), runtime_activity = Neither, broken = Neither, splat = false
+    )
     return TestCase(d, value, name, runtime_activity, broken, splat)
 end
 
 """
 Test Enzyme.gradient, both Forward and Reverse mode, against FiniteDifferences.grad.
 """
-function test_grad(case::TestCase; rtol=1e-6, atol=1e-6)
+function test_grad(case::TestCase; rtol = 1.0e-6, atol = 1.0e-6)
     @nospecialize
     f = case.func
     # We'll call the function as f(x...), so wrap in a singleton tuple if need be.
@@ -212,10 +212,10 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
         TestCase(LogUniform(0.15, 7.8), 7.1),
         TestCase(LogUniform(2.0, 3.0), 2.1),
         # TODO Broken tests, see https://github.com/EnzymeAD/Enzyme.jl/issues/1620
-        TestCase(NoncentralBeta(1.1, 1.1, 1.2), 0.8; broken=Both), # foreigncall (Rmath.dnbeta).
-        TestCase(NoncentralChisq(2, 3.0), 10.0; broken=Both), # foreigncall (Rmath.dnchisq).
-        TestCase(NoncentralF(2, 3, 1.1), 4.1; broken=Both), # foreigncall (Rmath.dnf).
-        TestCase(NoncentralT(1.3, 1.1), 0.1; broken=Both), # foreigncall (Rmath.dnt).
+        TestCase(NoncentralBeta(1.1, 1.1, 1.2), 0.8; broken = Both), # foreigncall (Rmath.dnbeta).
+        TestCase(NoncentralChisq(2, 3.0), 10.0; broken = Both), # foreigncall (Rmath.dnchisq).
+        TestCase(NoncentralF(2, 3, 1.1), 4.1; broken = Both), # foreigncall (Rmath.dnf).
+        TestCase(NoncentralT(1.3, 1.1), 0.1; broken = Both), # foreigncall (Rmath.dnt).
         TestCase(Normal(), 0.1),
         TestCase(Normal(0.0, 1.0), 1.0),
         TestCase(Normal(0.5, 1.0), 0.05),
@@ -223,7 +223,7 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
         TestCase(Normal(-0.1, 0.9), -0.3),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1603
         # foreigncall -- https://github.com/JuliaMath/SpecialFunctions.jl/blob/be1fa06fee58ec019a28fb0cd2b847ca83a5af9a/src/bessel.jl#L265
-        TestCase(NormalInverseGaussian(0.0, 1.0, 0.2, 0.1), 0.1; broken=Both),
+        TestCase(NormalInverseGaussian(0.0, 1.0, 0.2, 0.1), 0.1; broken = Both),
         TestCase(Pareto(1.0, 1.0), 3.5),
         TestCase(Pareto(1.1, 0.9), 3.1),
         TestCase(Pareto(1.0, 1.0), 1.4),
@@ -234,7 +234,7 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
         TestCase(Rayleigh(0.9), 1.1),
         TestCase(Rayleigh(0.55), 0.63),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1620
-        TestCase(Rician(0.5, 1.0), 2.1; broken=Both),  # foreigncall (Rmath.dnchisq).
+        TestCase(Rician(0.5, 1.0), 2.1; broken = Both),  # foreigncall (Rmath.dnchisq).
         TestCase(Semicircle(1.0), 0.9),
         TestCase(Semicircle(5.1), 5.05),
         TestCase(Semicircle(0.5), -0.1),
@@ -284,22 +284,22 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
         TestCase(MvNormal([0.2, 0.3], Symmetric(Diagonal([0.5, 0.4]))), [-0.1, 0.05]),
         TestCase(MvNormal([0.2, 0.3], Diagonal([0.5, 0.4])), [-0.1, 0.05]),
         # TODO Broken tests, see https://github.com/EnzymeAD/Enzyme.jl/issues/1991
-        TestCase(MvNormal([-0.15], _pdmat([1.1]')), [-0.05]; broken=Forward),
+        TestCase(MvNormal([-0.15], _pdmat([1.1]')), [-0.05]; broken = Forward),
         TestCase(
             MvNormal([0.2, -0.15], _pdmat([1.0 0.9; 0.7 1.1])), [0.05, -0.05];
-            broken=Forward
+            broken = Forward
         ),
         TestCase(MvNormal([0.2, -0.3], [0.5, 0.6]), [0.4, -0.3]),
         TestCase(MvNormalCanon([0.1, -0.1], _pdmat([0.5 0.4; 0.45 1.0])), [0.2, -0.25]),
         # TODO Broken tests, see https://github.com/EnzymeAD/Enzyme.jl/issues/1991
         TestCase(
             MvLogNormal(MvNormal([0.2, -0.1], _pdmat([1.0 0.9; 0.7 1.1]))), [0.5, 0.1];
-            broken=Forward
+            broken = Forward
         ),
         TestCase(product_distribution([Normal()]), [0.3]),
         TestCase(
             product_distribution([Normal(), Uniform()]), [-0.4, 0.3];
-            runtime_activity=Both
+            runtime_activity = Both
         ),
 
         #
@@ -316,12 +316,12 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
         TestCase(
             Wishart(5, _pdmat(randn(rng, 3, 3))),
             Symmetric(collect(_pdmat(randn(rng, 3, 3))));
-            broken=Forward
+            broken = Forward
         ),
         TestCase(
             InverseWishart(5, _pdmat(randn(rng, 3, 3))),
             Symmetric(collect(_pdmat(randn(rng, 3, 3))));
-            broken=Forward
+            broken = Forward
         ),
         # TODO Broken tests, see https://github.com/EnzymeAD/Enzyme.jl/issues/1820
         TestCase(
@@ -332,15 +332,15 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
                 _pdmat(randn(rng, 3, 3)),
             ),
             randn(rng, 2, 3);
-            broken=Both
+            broken = Both
         ),
-        TestCase(MatrixBeta(5, 6.0, 7.0), rand(rng, MatrixBeta(5, 6.0, 6.0)); broken=Both),
+        TestCase(MatrixBeta(5, 6.0, 7.0), rand(rng, MatrixBeta(5, 6.0, 6.0)); broken = Both),
         TestCase(
             MatrixFDist(6.0, 7.0, _pdmat(randn(rng, 5, 5))),
             rand(rng, MatrixFDist(6.0, 7.0, _pdmat(randn(rng, 5, 5))));
-            broken=Both
+            broken = Both
         ),
-        TestCase(LKJ(5, 1.1), rand(rng, LKJ(5, 1.1)); broken=Both),
+        TestCase(LKJ(5, 1.1), rand(rng, LKJ(5, 1.1)); broken = Both),
 
         #
         # Miscellaneous others
@@ -348,48 +348,48 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
 
         TestCase(
             (a, b, x) -> logpdf(InverseGamma(a, b), x), (1.5, 1.4, 0.4);
-            name="InverseGamma", splat=true
+            name = "InverseGamma", splat = true
         ),
         TestCase(
             (m, s, x) -> logpdf(NormalCanon(m, s), x), (0.1, 1.0, -0.5);
-            name="NormalCanon", splat=true
+            name = "NormalCanon", splat = true
         ),
-        TestCase(x -> logpdf(Categorical(x, 1 - x), 1), 0.3; name="Categorical"),
+        TestCase(x -> logpdf(Categorical(x, 1 - x), 1), 0.3; name = "Categorical"),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1995
         TestCase(
             (m, S, x) -> logpdf(MvLogitNormal(m, S), vcat(x, 1 - sum(x))),
             ([0.4, 0.6], _pdmat([0.9 0.4; 0.5 1.1]), [0.27, 0.24]);
-            name="MvLogitNormal", runtime_activity=Forward, broken=Forward, splat=true,
+            name = "MvLogitNormal", runtime_activity = Forward, broken = Forward, splat = true,
         ),
         TestCase(
             (a, b, α, β, x) -> logpdf(truncated(Beta(α, β), a, b), x),
             (0.1, 0.9, 1.1, 1.3, 0.4);
-            name="truncated Beta", splat=true
+            name = "truncated Beta", splat = true
         ),
         TestCase(
             (a, b, x) -> logpdf(truncated(Normal(), a, b), x),
             (-0.3, 0.3, 0.1);
-            name="allocs Normal", splat=true
+            name = "allocs Normal", splat = true
         ),
         TestCase(
             (a, b, α, β, x) -> logpdf(truncated(Uniform(α, β), a, b), x),
             (0.1, 0.9, -0.1, 1.1, 0.4);
-            name="allocs Uniform", splat=true
+            name = "allocs Uniform", splat = true
         ),
         TestCase(
             (a, x) -> logpdf(Dirichlet(a), [x, 1 - x]), ([1.5, 1.1], 0.6);
-            name="Dirichlet", splat=true, runtime_activity=Forward
+            name = "Dirichlet", splat = true, runtime_activity = Forward
         ),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1997
         TestCase(
             x -> logpdf(reshape(product_distribution([Normal(), Uniform()]), 1, 2), x),
             [2.1 0.7];
-            name="reshape"
+            name = "reshape"
         ),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1820
         TestCase(
             x -> logpdf(vec(LKJ(2, 1.1)), x), [1.0, 0.489, 0.489, 1.0];
-            name="vec", broken=Both
+            name = "vec", broken = Both
         ),
         # TODO Broken test, see https://github.com/EnzymeAD/Enzyme.jl/issues/1994
         TestCase(
@@ -403,7 +403,7 @@ _pdmat(A) = PDMat(_sym(A) + 5I)
                 return logpdf(LKJCholesky(2, v), C)
             end,
             (randn(rng, 2, 2), 1.1);
-            name="LKJCholesky", broken=Forward
+            name = "LKJCholesky", broken = Forward
         ),
     ]
 

[vchuravy]

julia: /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Instructions.h:1190: void llvm::ICmpInst::AssertOK(): Assertion `getOperand(0)->getType() == getOperand(1)->getType() && "Both operands to ICmp instruction are not of the same type!"' failed.

[3341] signal (6.-6): Aborted
in expression starting at /home/runner/work/Enzyme.jl/Enzyme.jl/test/integration/Distributions/runtests.jl:123
pthread_kill at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
gsignal at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
abort at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x7f251b22881a)
__assert_fail at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
AssertOK at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Instructions.h:1190 [inlined]
ICmpInst at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Instructions.h:1245 [inlined]
CreateICmp at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/IRBuilder.h:2180
CreateICmpEQ at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/IRBuilder.h:2067 [inlined]
handle_trtrs at /workspace/srcdir/Enzyme/build/Enzyme/BlasDerivatives.inc:15793
handleBLAS at /workspace/srcdir/Enzyme/build/Enzyme/BlasDerivatives.inc:45

[wsmoses]

@mhauru we recently got more CI resources, so I think we can start getting some of these merged. One caveat is that integration tests now need a bit more setup on github actions side to make them run on the new resources.

Would you be able to integrate something like this (https://github.com/EnzymeAD/Reactant.jl/pull/1584/files) into this PR (and other turing/etc integration PRs as well). Let's get these merged!

[wsmoses]

@penelopeysm after the DI integration tests land, can you help get this landable?

[penelopeysm]

Sure, happy to. Feel free to close this and I'll open a new PR against main with the same contents.

[wsmoses]

closing and @penelopeysm will open a new one!