Explicit Enzyme rules on Enzyme 0.13 (#350)

* Bump Enzyme to v0.13 * Mark more broken Enzyme tests * Bump Julia compat entry to 1.10 (#342) * Bump minimum Julia version to 1.10 * Use 'min' in CI * Bump versions of GHA (julia-actions needs to be v2 for 'min') * Update Project.toml --------- Co-authored-by: Hong Ge <[email protected]> * Mark Enzyme tests as unbroken * Explicit Enzyme rules * Mark a few tests as broken * Remove batch reverse mode tests completely * Generic Exception? * Skip failing tests * Simplify Enzyme tests --------- Co-authored-by: Markus Hauru <[email protected]> Co-authored-by: Penelope Yong <[email protected]> Co-authored-by: Hong Ge <[email protected]>
TuringLang · Nov 28, 2024 · 1480e79 · 1480e79
1 parent 9a879c1
commit 1480e79
Show file tree

Hide file tree

Showing 22 changed files with 486 additions and 454 deletions.
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
@@ -0,0 +1,7 @@
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "monthly"
diff --git a/.github/workflows/AD.yml b/.github/workflows/AD.yml
@@ -6,42 +6,39 @@ on:
       - master
   pull_request:
 
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
 jobs:
   test:
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - 'min'
+          - 'lts'
           - '1'
         os:
           - ubuntu-latest
           - macOS-latest
-        arch:
-          - x64
         AD:
           - Enzyme
           - ForwardDiff
           - Mooncake
           - 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 }}
-      - uses: julia-actions/julia-buildpkg@latest
-      - uses: julia-actions/julia-runtest@latest
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: AD
           AD: ${{ matrix.AD }}
diff --git a/.github/workflows/Docs.yml b/.github/workflows/Docs.yml
@@ -18,8 +18,8 @@ jobs:
   docs:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@latest
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
           version: '1'
       - name: Install dependencies

diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml
@@ -20,8 +20,8 @@ jobs:
   format:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@latest
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
           version: 1
       - name: Format code

diff --git a/.github/workflows/Interface.yml b/.github/workflows/Interface.yml
@@ -7,27 +7,31 @@ on:
       - master
   pull_request:
 
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
 jobs:
   test:
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - 'min'
+          - 'lts'
           - '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 }}
-      - uses: julia-actions/julia-buildpkg@latest
-      - uses: julia-actions/julia-runtest@latest
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: Interface
diff --git a/Project.toml b/Project.toml
@@ -4,10 +4,8 @@ version = "0.14.2"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
-ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -18,14 +16,12 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
@@ -36,7 +32,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
-BijectorsEnzymeExt = ["Enzyme", "EnzymeCore"]
+BijectorsEnzymeCoreExt = "EnzymeCore"
 BijectorsForwardDiffExt = "ForwardDiff"
 BijectorsLazyArraysExt = "LazyArrays"
 BijectorsMooncakeExt = "Mooncake"
@@ -46,15 +42,12 @@ BijectorsZygoteExt = "Zygote"
 
 [compat]
 ArgCheck = "1, 2"
-ChainRules = "1"
 ChainRulesCore = "0.10.11, 1"
 ChangesOfVariables = "0.1"
-Compat = "3.46, 4.2"
 Distributions = "0.25.33"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.9"
-Enzyme = "0.12.22"
-EnzymeCore = "0.7.8"
+EnzymeCore = "0.8.4"
 ForwardDiff = "0.10"
 Functors = "0.1, 0.2, 0.3, 0.4, 0.5"
 InverseFunctions = "0.1"
@@ -64,17 +57,15 @@ LogExpFunctions = "0.3.3"
 MappedArrays = "0.2.2, 0.3, 0.4"
 Mooncake = "0.4.19"
 Reexport = "0.2, 1"
-Requires = "0.5, 1"
 ReverseDiff = "1"
 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"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"

diff --git a/ext/BijectorsDistributionsADExt.jl b/ext/BijectorsDistributionsADExt.jl
@@ -1,38 +1,20 @@
 module BijectorsDistributionsADExt
 
-if isdefined(Base, :get_extension)
-    using Bijectors
-    using Bijectors: LinearAlgebra
-    using Bijectors.Distributions: AbstractMvLogNormal
-    using DistributionsAD:
-        TuringDirichlet,
-        TuringWishart,
-        TuringInverseWishart,
-        FillVectorOfUnivariate,
-        FillMatrixOfUnivariate,
-        MatrixOfUnivariate,
-        FillVectorOfMultivariate,
-        VectorOfMultivariate,
-        TuringScalMvNormal,
-        TuringDiagMvNormal,
-        TuringDenseMvNormal
-else
-    using ..Bijectors
-    using ..Bijectors: LinearAlgebra
-    using ..Bijectors.Distributions: AbstractMvLogNormal
-    using ..DistributionsAD:
-        TuringDirichlet,
-        TuringWishart,
-        TuringInverseWishart,
-        FillVectorOfUnivariate,
-        FillMatrixOfUnivariate,
-        MatrixOfUnivariate,
-        FillVectorOfMultivariate,
-        VectorOfMultivariate,
-        TuringScalMvNormal,
-        TuringDiagMvNormal,
-        TuringDenseMvNormal
-end
+using Bijectors
+using Bijectors: LinearAlgebra
+using Bijectors.Distributions: AbstractMvLogNormal
+using DistributionsAD:
+    TuringDirichlet,
+    TuringWishart,
+    TuringInverseWishart,
+    FillVectorOfUnivariate,
+    FillMatrixOfUnivariate,
+    MatrixOfUnivariate,
+    FillVectorOfMultivariate,
+    VectorOfMultivariate,
+    TuringScalMvNormal,
+    TuringDiagMvNormal,
+    TuringDenseMvNormal
 
 # Bijectors