Add DI integration tests

.github/workflows/CI.yml

@@ -230,34 +230,6 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false  # or true if you want CI to fail when Codecov fails
-  integration:
-    timeout-minutes: 20
-    name: Integration Tests - ${{ matrix.test }}
-    runs-on: ${{ matrix.os }}
-    env:
-      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - '1.10'
-        os:
-          - ubuntu-latest
-        test:
-          - DynamicExpressions
-          - Bijectors
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - name: "Run tests"
-        run: |
-            julia --color=yes --project=test/integration/${{ matrix.test }} -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")]); Pkg.instantiate()'
-            julia --color=yes --project=test/integration/${{ matrix.test }} --threads=auto --check-bounds=yes test/integration/${{ matrix.test }}/runtests.jl
-        shell: bash
   docs:
     timeout-minutes: 20
     name: Documentation

.github/workflows/Integration.yml

@@ -0,0 +1,56 @@
+name: Integration
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      - release-*
+    tags: '*'
+
+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:
+  integration:
+    timeout-minutes: 45
+    name: Integration Tests - ${{ matrix.package }} - Julia ${{ matrix.version }}
+    runs-on: ${{ matrix.os }}
+    container:
+      image: ${{ (contains(matrix.os, 'linux') && 'ghcr.io/enzymead/reactant-docker-images@sha256:91e1edb7a7c869d5a70db06e417f22907be0e67ca86641d48adcea221fedc674' ) || '' }}
+    env:
+      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+          - '1.11'
+        os:
+          - linux-x86-n2-32
+        package:
+          - Bijectors
+          - DifferentiationInterface
+          - DynamicExpressions
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - name: "Install Dependencies"
+        run: |
+            julia --color=yes --project=test/integration/${{ matrix.package }} --threads=auto --check-bounds=yes -O1 -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")])'
+        shell: bash
+        if: ${{ matrix.version == '1.10' }}
+      - name: "Instantiate"
+        run: |
+          julia --color=yes --project=test/integration/${{ matrix.package }} --threads=auto --check-bounds=yes -O1 -e 'using Pkg; Pkg.instantiate()'
+        shell: bash
+      - name: "Run tests"
+        run: |
+            julia --color=yes --project=test/integration/${{ matrix.package }} --threads=auto --check-bounds=yes -O1 test/integration/${{ matrix.package }}/runtests.jl
+        shell: bash

.gitignore

@@ -8,5 +8,5 @@ lib/EnzymeCore/Manifest.toml
 /docs/Manifest.toml
 /docs/build/
 /docs/src/generated/
-
+test/integration/**/Manifest.toml
 .vscode

test/integration/Bijectors/Project.toml

@@ -1,8 +1,14 @@
 [deps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
+[sources]
+Enzyme = { path = "../../.." }
+EnzymeCore = { path = "../../../lib/EnzymeCore" }
+
 [compat]
 Bijectors = "=0.13.16"
 FiniteDifferences = "0.12.32"

test/integration/Bijectors/runtests.jl

@@ -131,7 +131,11 @@ end
         sum_b_binv_test_case(Bijectors.VecCholeskyBijector(:U), 3),
         sum_b_binv_test_case(Bijectors.VecCholeskyBijector(:U), 0),
         sum_b_binv_test_case(Bijectors.Coupling(Bijectors.Shift, Bijectors.PartitionMask(3, [1], [2])), 3),
-        sum_b_binv_test_case(Bijectors.InvertibleBatchNorm(3), (3, 3)),
+        sum_b_binv_test_case(
+            Bijectors.InvertibleBatchNorm(3),
+            (3, 3),
+            broken = (VERSION >= v"1.11" ? Both : Neither)
+        ),
         sum_b_binv_test_case(Bijectors.LeakyReLU(0.2), 3),
         sum_b_binv_test_case(Bijectors.Logit(0.1, 0.3), 3),
         sum_b_binv_test_case(Bijectors.PDBijector(), (3, 3)),
@@ -186,6 +190,7 @@ end
             end,
             randn(rng, 7);
             name="PlanarLayer7",
+            broken = (VERSION >= v"1.11" ? Forward : Neither),
         ),
 
         TestCase(
@@ -197,6 +202,7 @@ end
             end,
             randn(rng, 11);
             name="PlanarLayer11",
+            broken = (VERSION >= v"1.11" ? Forward : Neither),
         ),
     ]
 

test/integration/DifferentiationInterface/Project.toml

@@ -0,0 +1,18 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Enzyme = { path = "../../.." }
+EnzymeCore = { path = "../../../lib/EnzymeCore" }
+
+[compat]
+ADTypes = "1.17.0"
+DifferentiationInterface = "=0.7.7"
+DifferentiationInterfaceTest = "=0.10.2"
+StaticArrays = "1.9"

test/integration/DifferentiationInterface/README.md

@@ -0,0 +1,35 @@
+# DI integration tests for Enzyme
+
+This folder contains tests to ensure that changes to Enzyme do not break integration with [DifferentiationInterface](https://github.com/JuliaDiff/DifferentiationInterface.jl) (DI).
+
+## Relevant source files
+
+The test utilities used here come from the sibling package [DifferentiationInterfaceTest](https://github.com/JuliaDiff/DifferentiationInterface.jl/tree/main/DifferentiationInterfaceTest) (DIT).
+Correctness checking itself is implemented in [`src/tests/correctness_eval.jl`](https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/ed5655a90bf9f3a6092904070d353a9d705ebdc4/DifferentiationInterfaceTest/src/tests/correctness_eval.jl), which is where you will see test errors originating.
+Test scenarios are located in [`src/scenarios`](https://github.com/JuliaDiff/DifferentiationInterface.jl/tree/ed5655a90bf9f3a6092904070d353a9d705ebdc4/DifferentiationInterfaceTest/src/scenarios) (especially [`src/scenarios/default.jl`](https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/ed5655a90bf9f3a6092904070d353a9d705ebdc4/DifferentiationInterfaceTest/src/scenarios/default.jl) and [`src/scenarios/modify.jl`](https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/ed5655a90bf9f3a6092904070d353a9d705ebdc4/DifferentiationInterfaceTest/src/scenarios/modify.jl)) and in package extensions.
+The structure of a `Scenario` is defined in [`src/scenarios/scenario.jl`](https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/ed5655a90bf9f3a6092904070d353a9d705ebdc4/DifferentiationInterfaceTest/src/scenarios/scenario.jl).
+
+Scenario generation relies on internals of DifferentiationInterfaceTest, which is why its version is pinned in the `Project.toml`.
+
+## Interpreting test errors
+
+The most common test error you will see looks like
+
+```julia
+Correctness: Test Failed at .../src/tests/correctness_eval.jl:...
+  Expression: res1_out1_noval ≈ scen.res1
+```
+
+Each test scenario `scen` contains a first-order result `res1` and a second-order result `res2`, which are the reference values we compare our autodiff results (i.e. the output of `DI.gradient` or `DI.jacobian`) against.
+The suffixes of the left-hand term are defined as follows:
+
+- `_out` for the output of the operator, `_in` for the input if it is in-place
+- `1` for the first call to the operator, `2` for the second call (which is used to check that the preparation object has not been altered and can safely be reused)
+- `_val` if the operator also returns the value of the function (like `DI.value_and_gradient`), `_noval` otherwise
+
+As you can see, several variants of each operator are tested, so a single bug will give rise to many different errors. In addition, different preparation mechanisms are also cycled through.
+The testset summary at the end of the CI log is probably the right place to start hunting down issues.
+
+## What to do if a test fails
+
+Open an issue on the DI repo with a link to the relevant PR or CI log.

test/integration/DifferentiationInterface/runtests.jl

@@ -0,0 +1,86 @@
+using ADTypes: AutoEnzyme
+using DifferentiationInterface: DifferentiationInterface
+using DifferentiationInterfaceTest:
+    default_scenarios,
+    sparse_scenarios,
+    static_scenarios,
+    test_differentiation,
+    function_place,
+    operator_place,
+    FIRST_ORDER,
+    SECOND_ORDER,
+    Scenario
+using Enzyme: Enzyme
+using EnzymeCore: Forward, Reverse, Const, Duplicated
+using StaticArrays: StaticArrays
+using Test
+
+logging = get(ENV, "CI", "false") == "false"
+
+function remove_matrices(scens::Vector{<:Scenario})  # TODO: remove
+    return filter(s -> s.x isa Union{Number, AbstractVector} && s.y isa Union{Number, AbstractVector}, scens)
+end
+
+backends = [
+    AutoEnzyme(; function_annotation = Const),
+    AutoEnzyme(; mode = Forward),
+    AutoEnzyme(; mode = Reverse),
+]
+
+duplicated_backends = [
+    AutoEnzyme(; mode = Forward, function_annotation = Duplicated),
+    AutoEnzyme(; mode = Reverse, function_annotation = Duplicated),
+]
+
+@testset verbose = true "DifferentiationInterface integration" begin
+    test_differentiation(
+        backends,
+        default_scenarios(; include_constantified = true);
+        excluded = SECOND_ORDER,
+        logging,
+        testset_name = "Generic first order",
+    )
+
+    test_differentiation(
+        backends[1],
+        remove_matrices(default_scenarios(; include_constantified = true));
+        excluded = FIRST_ORDER,
+        logging,
+        testset_name = "Generic second order",
+    )
+
+    test_differentiation(
+        backends[2],
+        remove_matrices(
+            default_scenarios(;
+                include_normal = false,
+                include_cachified = true,
+                include_constantorcachified = true,
+                use_tuples = true,
+            )
+        );
+        excluded = SECOND_ORDER,
+        logging,
+        testset_name = "Caches",
+    )
+
+    test_differentiation(
+        duplicated_backends,
+        remove_matrices(default_scenarios(; include_normal = false, include_closurified = true));
+        excluded = SECOND_ORDER,
+        logging,
+        testset_name = "Closures",
+    )
+
+    filtered_static_scenarios = filter(remove_matrices(static_scenarios())) do s
+        operator_place(s) == :out && function_place(s) == :out
+    end
+
+    test_differentiation(
+        backends[2:3],
+        filtered_static_scenarios;
+        excluded = SECOND_ORDER,
+        logging,
+        testset_name = "Static arrays",
+    )
+end

test/integration/DynamicExpressions/Project.toml

@@ -1,5 +1,11 @@
 [deps]
 DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[sources]
+Enzyme = { path = "../../.." }
+EnzymeCore = { path = "../../../lib/EnzymeCore" }
 
 [compat]
 DynamicExpressions = "=0.18.5"