Refactor CI

.github/workflows/AD.yml → .github/workflows/CI.yml

@@ -1,4 +1,4 @@
-name: AD tests
+name: CI
 
 on:
   push:
@@ -14,33 +14,30 @@ concurrency:
 
 jobs:
   test:
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ matrix.runner.os }}
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - 'min'
-          - '1'
-        os:
-          - ubuntu-latest
-          - macOS-latest
-        AD:
-          - Enzyme
-          - ForwardDiff
-          - Mooncake
-          - Tracker
-          - ReverseDiff
+        runner:
+          - version: '1'
+            os: 'ubuntu-latest'
+          - version: '1'
+            os: 'macos-latest'
+          - version: '1'
+            os: 'windows-latest'
+          - version: 'min'
+            os: 'ubuntu-latest'
+        group:
+          - 'Interface'
+          - 'AD'
+
     steps:
       - uses: actions/checkout@v5
-
       - uses: julia-actions/setup-julia@v2
         with:
-          version: ${{ matrix.version }}
-
+          version: ${{ matrix.runner.version }}
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
-
       - uses: julia-actions/julia-runtest@v1
-        continue-on-error: ${{ matrix.AD == 'Enzyme' && matrix.version == '1' }}
         env:
-          GROUP: AD
-          AD: ${{ matrix.AD }}
+          GROUP: ${{ matrix.group }}

.github/workflows/DocTests.yml

@@ -0,0 +1,41 @@
+# We want to only run doctests on a single version of Julia, because
+# things like error messages / output can change between versions and
+# is fragile to test against.
+name: Doctests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  merge_group:
+    types: [checks_requested]
+
+# needed to allow julia-actions/cache to delete old caches that it has created
+permissions:
+  actions: write
+  contents: read
+
+# Cancel existing tests on the same PR if a new commit is added to a pull request
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
+
+      - uses: julia-actions/cache@v2
+
+      - uses: julia-actions/julia-buildpkg@v1
+
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          GROUP: Doctests

docs/Project.toml

@@ -5,6 +5,6 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 
 [compat]
-Documenter = "1.14"
+Documenter = "1"
 Functors = "0.3"
 StableRNGs = "1"

docs/make.jl

@@ -15,4 +15,5 @@ makedocs(;
         "Examples" => "examples.md",
     ],
     checkdocs=:exports,
+    doctest=false,
 )

test/Project.toml

@@ -4,6 +4,7 @@ AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -32,6 +33,7 @@ AdvancedHMC = "0.6, 0.7, 0.8"
 ChainRulesTestUtils = "0.7, 1"
 ChangesOfVariables = "0.1"
 Combinatorics = "1.0.2"
+DifferentiationInterface = "0.7.7"
 DistributionsAD = "0.6.3"
 Documenter = "1"
 Enzyme = "0.13.12"

test/ad/chainrules.jl

@@ -1,6 +1,11 @@
-using Random: Xoshiro
+module BijectorsChainRulesTests
+
+using Bijectors
 using LinearAlgebra
-using ChainRulesTestUtils: ChainRulesCore
+using ChainRulesTestUtils: ChainRulesCore, test_frule, test_rrule, ⊢, ChainRulesTestUtils
+using FiniteDifferences: FiniteDifferences
+using Random: Xoshiro
+using Test
 
 # HACK: This is a workaround to test `Bijectors._inv_link_chol_lkj` which produces an
 # upper-triangular `Matrix`, leading to `test_rrule` comaring the _full_ `Matrix`,
@@ -29,42 +34,6 @@ end
     test_frule(Bijectors.find_alpha, x, y, z)
     test_rrule(Bijectors.find_alpha, x, y, z)
 
-    if @isdefined Mooncake
-        rng = Xoshiro(123456)
-        @testset "$mode" for mode in (Mooncake.ReverseMode, Mooncake.ForwardMode)
-            Mooncake.TestUtils.test_rule(
-                rng,
-                Bijectors.find_alpha,
-                x,
-                y,
-                z;
-                is_primitive=true,
-                perf_flag=:none,
-                mode=mode,
-            )
-            Mooncake.TestUtils.test_rule(
-                rng,
-                Bijectors.find_alpha,
-                x,
-                y,
-                3;
-                is_primitive=true,
-                perf_flag=:none,
-                mode=mode,
-            )
-            Mooncake.TestUtils.test_rule(
-                rng,
-                Bijectors.find_alpha,
-                x,
-                y,
-                UInt32(3);
-                is_primitive=true,
-                perf_flag=:none,
-                mode=mode,
-            )
-        end
-    end
-
     test_rrule(
         Bijectors.combine,
         Bijectors.PartitionMask(3, [1], [2]) ⊢ ChainRulesTestUtils.NoTangent(),
@@ -182,3 +151,5 @@ end
         end
     end
 end
+
+end # module BijectorsChainRulesTests

test/ad/corr.jl

@@ -1,35 +1,49 @@
-@testset "AD for VecCorrBijector" begin
-    d = 4
-    dist = LKJ(d, 2.0)
-    b = bijector(dist)
-    binv = inverse(b)
+using Enzyme: ForwardMode
 
-    x = rand(dist)
-    y = b(x)
+@testset "VecCorrBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES
+    ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:Enzyme.ForwardMode}
 
-    test_ad(y) do x
-        sum(transform(b, binv(x)))
-    end
+    @testset "d = $d" for d in (1, 2, 4)
+        dist = LKJ(d, 2.0)
+        b = bijector(dist)
+        binv = inverse(b)
+
+        x = rand(dist)
+        y = b(x)
 
-    test_ad(y) do y
-        sum(transform(binv, y))
+        roundtrip(y) = sum(transform(b, binv(y)))
+        inverse_only(y) = sum(transform(binv, y))
+        if d == 4 && ENZYME_FWD_AND_1p11
+            @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad(
+                roundtrip, adtype, y
+            )
+            @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad(
+                inverse_only, adtype, y
+            )
+        else
+            test_ad(roundtrip, adtype, y)
+            test_ad(inverse_only, adtype, y)
+        end
     end
 end
 
-@testset "AD for VecCholeskyBijector" begin
-    d = 4
-    dist = LKJCholesky(d, 2.0)
-    b = bijector(dist)
-    binv = inverse(b)
+@testset "VecCholeskyBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES
+    @testset "d = $d, uplo = $uplo" for d in (1, 2, 4), uplo in ('U', 'L')
+        dist = LKJCholesky(d, 2.0, uplo)
+        b = bijector(dist)
+        binv = inverse(b)
 
-    x = rand(dist)
-    y = b(x)
+        x = rand(dist)
+        y = b(x)
+        cholesky_to_triangular =
+            uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower
 
-    test_ad(y) do y
-        sum(transform(b, binv(y)))
-    end
+        roundtrip(y) = sum(transform(b, binv(y)))
+        test_ad(roundtrip, adtype, y)
 
-    test_ad(y) do y
-        sum(Bijectors.cholesky_upper(transform(binv, y)))
+        # we need to tack on `cholesky_upper`/`cholesky_lower`, because directly calling
+        # `sum` on a LinearAlgebra.Cholesky doesn't give a scalar
+        inverse_only(y) = sum(cholesky_to_triangular(transform(binv, y)))
+        test_ad(inverse_only, adtype, y)
     end
 end