Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 18 additions & 21 deletions .github/workflows/AD.yml → .github/workflows/CI.yml
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
name: AD tests
name: CI

on:
push:
Expand All @@ -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 }}
41 changes: 41 additions & 0 deletions .github/workflows/DocTests.yml
Original file line number Diff line number Diff line change
@@ -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
36 changes: 0 additions & 36 deletions .github/workflows/Interface.yml

This file was deleted.

2 changes: 1 addition & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
3 changes: 3 additions & 0 deletions docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,7 @@ makedocs(;
"Examples" => "examples.md",
],
checkdocs=:exports,
doctest=false,
# TODO: Remove this
warnonly=true,
Comment on lines +19 to +20
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a replacement for strict=false which no longer works with Documenter.jl v1.0. It converts all docs errors into warnings. This is not good, but I didn't want fix this in this PR as well, it should be a separate PR.

)
2 changes: 2 additions & 0 deletions test/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down
47 changes: 9 additions & 38 deletions test/ad/chainrules.jl
Original file line number Diff line number Diff line change
@@ -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`,
Expand Down Expand Up @@ -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)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

moved to test/ad/mooncake.jl

@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(),
Expand Down Expand Up @@ -182,3 +151,5 @@ end
end
end
end

end # module BijectorsChainRulesTests
52 changes: 27 additions & 25 deletions test/ad/corr.jl
Original file line number Diff line number Diff line change
@@ -1,35 +1,37 @@
@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
@testset "d = $d" for d in (1, 2, 4)
dist = LKJ(d, 2.0)
b = bijector(dist)
binv = inverse(b)

test_ad(y) do x
sum(transform(b, binv(x)))
end
x = rand(dist)
y = b(x)

test_ad(y) do y
sum(transform(binv, y))
# roundtrip
test_ad(y -> sum(transform(b, binv(y))), adtype, y)
# inverse only
test_ad(y -> sum(transform(binv, y)), adtype, y)
end
end

@testset "AD for VecCholeskyBijector" begin
d = 4
dist = LKJCholesky(d, 2.0)
b = bijector(dist)
binv = inverse(b)

x = rand(dist)
y = b(x)
@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)

test_ad(y) do y
sum(transform(b, binv(y)))
end
x = rand(dist)
y = b(x)
cholesky_to_triangular =
uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower

test_ad(y) do y
sum(Bijectors.cholesky_upper(transform(binv, y)))
# roundtrip
test_ad(y -> sum(transform(b, binv(y))), adtype, y)
# inverse (we need to tack on `cholesky_upper`/`cholesky_lower`,
# because directly calling `sum` on a LinearAlgebra.Cholesky doesn't
# give a scalar)
test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y)
end
end
93 changes: 51 additions & 42 deletions test/ad/enzyme.jl
Original file line number Diff line number Diff line change
@@ -1,51 +1,60 @@
@testset "Enzyme: Bijectors.find_alpha" begin
x = randn()
y = expm1(randn())
z = randn()

@testset "forward" begin
# No batches
@testset for RT in (Const, Duplicated, DuplicatedNoNeed),
Tx in (Const, Duplicated),
Ty in (Const, Duplicated),
Tz in (Const, Duplicated)

# Rule not picked up by Enzyme on Julia 1.11?!
# Ref https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968
if VERSION >= v"1.11" && Tx <: Const && Ty <: Const && Tz <: Const
continue
module BijectorsEnzymeRulesTests

using Bijectors
using Enzyme
using EnzymeTestUtils: test_forward, test_reverse
using Test

# This entire test suite is broken on 1.11.
#
# https://github.com/EnzymeAD/Enzyme.jl/issues/2121
# https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968
#
# Ideally we'd use `@test_throws`. However, that doesn't work because
# `test_forward` itself calls `@test`, and the error is captured by that
# `@test`, not our `@test_throws`. Consequently `@test_throws` doesn't actually
# see any error. Weird Julia behaviour.

@static if VERSION < v"1.11"
@testset "Enzyme: Bijectors.find_alpha" begin
x = randn()
y = expm1(randn())
z = randn()

@testset "forward" begin
# No batches
@testset for RT in (Const, Duplicated, DuplicatedNoNeed),
Tx in (Const, Duplicated),
Ty in (Const, Duplicated),
Tz in (Const, Duplicated)

test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
end

test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
end

# Batches
@testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed),
Tx in (Const, BatchDuplicated),
Ty in (Const, BatchDuplicated),
Tz in (Const, BatchDuplicated)
# Batches
@testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed),
Tx in (Const, BatchDuplicated),
Ty in (Const, BatchDuplicated),
Tz in (Const, BatchDuplicated)

# Rule not picked up by Enzyme on Julia 1.11?!
# Ref https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968
if VERSION >= v"1.11" && Tx <: Const && Ty <: Const && Tz <: Const
continue
test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
end
end
@testset "reverse" begin
# No batches
@testset for RT in (Const, Active),
Tx in (Const, Active),
Ty in (Const, Active),
Tz in (Const, Active)

test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
end

test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
# TODO: Test batch mode
# This is a bit problematic since Enzyme does not support all combinations of activities currently
# https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728
end
end
@testset "reverse" begin
# No batches
@testset for RT in (Const, Active),
Tx in (Const, Active),
Ty in (Const, Active),
Tz in (Const, Active)

test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
end
end

# TODO: Test batch mode
# This is a bit problematic since Enzyme does not support all combinations of activities currently
# https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728
end
end
Loading
Loading