add ForwardDiff@1 (#378)

* add ForwardDiff@1

* Force ForwardDiff=1 in tests

* Don't test Dirichlet AD at non-differentiable points

* Don't run duplicate tests

* Use ForwardDiff 1.0.1, fix LKJCholesky Jacobian test

* Remove dead code

* Don't force FD=1 in test suite

* Add reference to ForwardDiff issue

Author: penelopeysm

Project.toml

@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.15.6"
+version = "0.15.7"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -51,7 +51,7 @@ Distributions = "0.25.33"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.8.4"
-ForwardDiff = "0.10"
+ForwardDiff = "0.10, 1.0.1"
 Functors = "0.1, 0.2, 0.3, 0.4, 0.5"
 InverseFunctions = "0.1"
 IrrationalConstants = "0.1, 0.2"

src/Bijectors.jl

@@ -86,7 +86,6 @@ _eps(::Type{<:Integer}) = eps(Float64)
 
 function _clamp(x, a, b)
     T = promote_type(typeof(x), typeof(a), typeof(b))
-    ϵ = _eps(T)
     clamped_x = ifelse(x < a, convert(T, a), ifelse(x > b, convert(T, b), x))
     DEBUG && _debug("x = $x, bounds = $((a, b)), clamped_x = $clamped_x")
     return clamped_x

test/Project.toml

@@ -39,7 +39,7 @@ Enzyme = "0.13.12"
 EnzymeTestUtils = "0.2.1"
 FillArrays = "1"
 FiniteDifferences = "0.11, 0.12"
-ForwardDiff = "0.10.12"
+ForwardDiff = "0.10, 1.0.1"
 Functors = "0.1, 0.2, 0.3, 0.4, 0.5"
 InverseFunctions = "0.1"
 LazyArrays = "1, 2"

test/interface.jl

@@ -145,12 +145,10 @@ end
 @testset "Multivariate" begin
     vector_dists = [
         Dirichlet(2, 3),
-        Dirichlet([1000 * one(Float64), eps(Float64)]),
-        Dirichlet([eps(Float64), 1000 * one(Float64)]),
+        Dirichlet([10.0, 0.1]),
+        Dirichlet([0.1, 10.0]),
         MvNormal(randn(10), Diagonal(exp.(randn(10)))),
         MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
-        Dirichlet([1000 * one(Float64), eps(Float64)]),
-        Dirichlet([eps(Float64), 1000 * one(Float64)]),
         MvTDist(1, randn(10), Matrix(Diagonal(exp.(randn(10))))),
         transformed(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
         transformed(MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10)))))),
@@ -173,15 +171,7 @@ end
             # similar to what we do in test/transform.jl for Dirichlet
             if dist isa Dirichlet
                 b = Bijectors.SimplexBijector()
-                # HACK(torfjelde): Calling `rand(dist)` will sometimes lead to `[0.999..., 0.0]`
-                # which in turn will lead to differences between `ForwardDiff.jacobian`
-                # and `logabsdetjac` due to how we handle the boundary values in `SimplexBijector`.
-                # We therefore test the realizations _on_ the boundary rather if we're near the boundary.
-                x = if any(rand(dist) .> 0.9999)
-                    [0.0, 1.0][sortperm(rand(dist))]
-                else
-                    rand(dist)
-                end
+                x = rand(dist)
                 y = b(x)
                 @test b(param(x)) isa TrackedArray
                 @test logabsdet(ForwardDiff.jacobian(b, x)[:, 1:(end - 1)])[1] ≈

test/transform.jl

@@ -153,32 +153,10 @@ end
             Dirichlet([eps(Float64), 1000 * one(Float64)]),
             MvNormal(randn(10), Diagonal(exp.(randn(10)))),
             MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
-            Dirichlet([1000 * one(Float64), eps(Float64)]),
-            Dirichlet([eps(Float64), 1000 * one(Float64)]),
         ]
         for dist in vector_dists
             if dist isa Dirichlet
                 single_sample_tests(dist)
-
-                # This should fail at the minute. Not sure what the correct way to test this is.
-
-                # Workaround for intermittent test failures, result of `logpdf_with_trans(dist, x, true)`
-                # is incorrect for `x == [0.9999999999999998, 0.0]`:
-                x =
-                    if params(dist) ==
-                        params(Dirichlet([1000 * one(Float64), eps(Float64)]))
-                        [1.0, 0.0]
-                    else
-                        rand(dist)
-                    end
-                # `Dirichlet` is no longer mapping between spaces of the same dimensionality,
-                # so the block below no longer works.
-                if !(dist isa Dirichlet)
-                    logpdf_turing = logpdf_with_trans(dist, x, true)
-                    J = ForwardDiff.jacobian(x -> link(dist, x), x)
-                    @test logpdf(dist, x .+ ϵ) - _logabsdet(J) ≈ logpdf_turing
-                end
-
                 # Issue #12
                 stepsize = 1e10
                 dim = Bijectors.output_length(bijector(dist), length(dist))
@@ -240,14 +218,15 @@ end
     @testset "uplo: $uplo" for uplo in [:L, :U]
         dist = LKJCholesky(3, 1, uplo)
         single_sample_tests(dist)
-
         x = rand(dist)
-
-        inds = [
-            LinearIndices(size(x))[I] for I in CartesianIndices(size(x)) if
-            (uplo === :L && I[2] < I[1]) || (uplo === :U && I[2] > I[1])
-        ]
         J = ForwardDiff.jacobian(z -> link(dist, Cholesky(z, x.uplo, x.info)), x.UL)
+        # Remove columns of Jacobian that are all zero (i.e. those
+        # corresponding to entries above the diagonal for uplo = :U, or below
+        # the diagonal for uplo = :L). This slightly unscientific approach
+        # based on filter() is needed to handle both ForwardDiff 0.10 and 1 as
+        # the exact indices will differ for the two versions; see
+        # https://github.com/JuliaDiff/ForwardDiff.jl/issues/738.
+        inds = filter(i -> !all(iszero, J[:, i]), 1:size(J, 2))
         J = J[:, inds]
         logpdf_turing = logpdf_with_trans(dist, x, true)
         @test logpdf(dist, x) - _logabsdet(J) ≈ logpdf_turing