Bump Enzyme to v0.13

.github/workflows/AD.yml

@@ -13,7 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - 'min'
           - '1'
         os:
           - ubuntu-latest
@@ -27,16 +27,9 @@ jobs:
           - Tracker
           - ReverseDiff
           - Zygote
-        exclude:
-          - version: 1.6
-            AD: Mooncake
-          # TODO(mhauru) Hopefully can enable Enzyme on older versions at some point, see
-          # discussion in https://github.com/TuringLang/Bijectors.jl/pull.
-          - version: 1.6
-            AD: Enzyme
     steps:
-      - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@v1
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}

.github/workflows/Interface.yml

@@ -14,16 +14,16 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - 'min'
           - '1'
         os:
           - ubuntu-latest
           - macOS-latest
         arch:
           - x64
     steps:
-      - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@v1
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}

Project.toml

@@ -36,7 +36,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
-BijectorsEnzymeExt = ["Enzyme", "EnzymeCore"]
+BijectorsEnzymeExt = "Enzyme"
 BijectorsForwardDiffExt = "ForwardDiff"
 BijectorsLazyArraysExt = "LazyArrays"
 BijectorsMooncakeExt = "Mooncake"
@@ -53,8 +53,7 @@ Compat = "3.46, 4.2"
 Distributions = "0.25.33"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.9"
-Enzyme = "0.12.22"
-EnzymeCore = "0.7.8"
+Enzyme = "0.13.12"
 ForwardDiff = "0.10"
 Functors = "0.1, 0.2, 0.3, 0.4, 0.5"
 InverseFunctions = "0.1"
@@ -70,7 +69,7 @@ Roots = "1.3.15, 2"
 Statistics = "1"
 Tracker = "0.2"
 Zygote = "0.6.63"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"

test/Project.toml

@@ -32,7 +32,7 @@ ChangesOfVariables = "0.1"
 Combinatorics = "1.0.2"
 Compat = "3.46, 4.2"
 DistributionsAD = "0.6.3"
-Enzyme = "0.12.22"
+Enzyme = "0.13.12"
 FillArrays = "1"
 FiniteDifferences = "0.11, 0.12"
 ForwardDiff = "0.10.12"
@@ -45,4 +45,4 @@ MCMCDiagnosticTools = "0.3"
 ReverseDiff = "1.4.2"
 Tracker = "0.2.11"
 Zygote = "0.6.63"
-julia = "1.3"
+julia = "1.10"

test/ad/utils.jl

@@ -12,6 +12,10 @@ function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6)
                 :Enzyme,
                 :EnzymeForward,
                 :EnzymeReverse,
+                # The `Crash` ones indicate that the error will cause a Julia crash, and
+                # thus we can't even run `@test_broken on it.
+                :EnzymeForwardCrash,
+                :EnzymeReverseCrash,
             )
         )
             error("Unknown broken AD backend: $b")
@@ -47,34 +51,39 @@ function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6)
         end
     end
 
-    # TODO(mhauru) The version bound should be relaxed once some Enzyme issues get
-    # sorted out. I think forward mode will remain broken for versions <= 1.6 due to
-    # some Julia bug. See https://github.com/EnzymeAD/Enzyme.jl/issues/1629 and
-    # discussion in https://github.com/TuringLang/Bijectors.jl/pull/318.
-    if (AD == "All" || AD == "Enzyme") && VERSION >= v"1.10"
+    if (AD == "All" || AD == "Enzyme")
         forward_broken = :EnzymeForward in broken || :Enzyme in broken
         reverse_broken = :EnzymeReverse in broken || :Enzyme in broken
-        if forward_broken
-            @test_broken(
-                collect(et, Enzyme.gradient(Enzyme.Forward, f, x)) ≈ finitediff,
-                rtol = rtol,
-                atol = atol
-            )
-        else
-            @test(
-                collect(et, Enzyme.gradient(Enzyme.Forward, f, x)) ≈ finitediff,
-                rtol = rtol,
-                atol = atol
-            )
+        if !(:EnzymeForwardCrash in broken)
+            if forward_broken
+                @test_broken(
+                    collect(et, Enzyme.gradient(Enzyme.Forward, f, x)[1]) ≈ finitediff,
+                    rtol = rtol,
+                    atol = atol
+                )
+            else
+                @test(
+                    collect(et, Enzyme.gradient(Enzyme.Forward, f, x)[1]) ≈ finitediff,
+                    rtol = rtol,
+                    atol = atol
+                )
+            end
         end
-        if reverse_broken
-            @test_broken(
-                Enzyme.gradient(Enzyme.Reverse, f, x) ≈ finitediff, rtol = rtol, atol = atol
-            )
-        else
-            @test(
-                Enzyme.gradient(Enzyme.Reverse, f, x) ≈ finitediff, rtol = rtol, atol = atol
-            )
+
+        if !(:EnzymeReverseCrash in broken)
+            if reverse_broken
+                @test_broken(
+                    collect(et, Enzyme.gradient(Enzyme.Reverse, f, x)[1]) ≈ finitediff,
+                    rtol = rtol,
+                    atol = atol
+                )
+            else
+                @test(
+                    collect(et, Enzyme.gradient(Enzyme.Reverse, f, x)[1]) ≈ finitediff,
+                    rtol = rtol,
+                    atol = atol
+                )
+            end
         end
     end