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 75.47%. Comparing base (037dfed) to head (ac42029).
Report is 355 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1819      +/-   ##
==========================================
+ Coverage   67.50%   75.47%   +7.96%     
==========================================
  Files          31       56      +25     
  Lines       12668    16742    +4074     
==========================================
+ Hits         8552    12636    +4084     
+ Misses       4116     4106      -10     

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

[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?

[yebai]

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

Maybe consider running these integration tests in a broader range of platforms by default.

[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
         ),
     ]