diff --git a/.github/workflows/JuliaNightly.yml b/.github/workflows/JuliaNightly.yml deleted file mode 100644 index 41642fa2..00000000 --- a/.github/workflows/JuliaNightly.yml +++ /dev/null @@ -1,33 +0,0 @@ -name: JuliaNightly - -on: - push: - branches: - - master - pull_request: - branches: - - master - -jobs: - test: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - uses: julia-actions/setup-julia@v1 - with: - version: 'nightly' - arch: x64 - - uses: actions/cache@v1 - env: - cache-name: cache-artifacts - with: - path: ~/.julia/artifacts - key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} - restore-keys: | - ${{ runner.os }}-test-${{ env.cache-name }}- - ${{ runner.os }}-test- - ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@latest - - uses: julia-actions/julia-runtest@latest - env: - GROUP: AdvancedVI diff --git a/README.md b/README.md index 08c8bec3..1ca0d848 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,6 @@ [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://turinglang.org/AdvancedVI.jl/stable/) [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://turinglang.org/AdvancedVI.jl/dev/) [![Build Status](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml/badge.svg?branch=master)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml?query=branch%3Amaster) -[![JuliaNightly](https://github.com/TuringLang/AdvancedVI.jl/workflows/JuliaNightly/badge.svg?branch=master)](https://github.com/TuringLang/AdvancedVI.jl/actions?query=workflow%3AJuliaNightly+branch%3Amaster) [![Coverage](https://codecov.io/gh/TuringLang/AdvancedVI.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/TuringLang/AdvancedVI.jl) # AdvancedVI.jl diff --git a/ext/AdvancedVIEnzymeExt.jl b/ext/AdvancedVIEnzymeExt.jl index 8333299f..29a0a695 100644 --- a/ext/AdvancedVIEnzymeExt.jl +++ b/ext/AdvancedVIEnzymeExt.jl @@ -11,15 +11,21 @@ else using ..AdvancedVI: ADTypes, DiffResults end -# Enzyme doesn't support f::Bijectors (see https://github.com/EnzymeAD/Enzyme.jl/issues/916) function AdvancedVI.value_and_gradient!( - ad::ADTypes.AutoEnzyme, f, θ::AbstractVector{T}, out::DiffResults.MutableDiffResult + ad::ADTypes.AutoEnzyme, + f, + θ::AbstractVector{T}, + out::DiffResults.MutableDiffResult, ) where {T<:Real} - y = f(θ) - DiffResults.value!(out, y) ∇θ = DiffResults.gradient(out) fill!(∇θ, zero(T)) - Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(θ, ∇θ)) + _, y = Enzyme.autodiff( + Enzyme.ReverseWithPrimal, + f, + Enzyme.Active, + Enzyme.Duplicated(θ, ∇θ), + ) + DiffResults.value!(out, y) return out end diff --git a/test/Project.toml b/test/Project.toml index a0dba17f..16370aee 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" @@ -26,6 +27,7 @@ ADTypes = "0.2.1, 1" Bijectors = "0.13" Distributions = "0.25.100" DistributionsAD = "0.6.45" +Enzyme = "0.12" FillArrays = "1.6.1" ForwardDiff = "0.10.36" Functors = "0.4.5" diff --git a/test/interface/ad.jl b/test/interface/ad.jl index be4ca34e..488bf4e3 100644 --- a/test/interface/ad.jl +++ b/test/interface/ad.jl @@ -3,11 +3,11 @@ using Test @testset "ad" begin @testset "$(adname)" for (adname, adsymbol) ∈ Dict( - :ForwardDiff => AutoForwardDiff(), - :ReverseDiff => AutoReverseDiff(), - :Zygote => AutoZygote(), - # :Enzyme => AutoEnzyme() # Currently not tested against - ) + :ForwardDiff => AutoForwardDiff(), + :ReverseDiff => AutoReverseDiff(), + :Zygote => AutoZygote(), + :Enzyme => AutoEnzyme(), + ) D = 10 A = randn(D, D) λ = randn(D) diff --git a/test/runtests.jl b/test/runtests.jl index 3bd13144..5cde6077 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -18,7 +18,7 @@ using DistributionsAD using LogDensityProblems using Optimisers using ADTypes -using ForwardDiff, ReverseDiff, Zygote +using Enzyme, ForwardDiff, ReverseDiff, Zygote using AdvancedVI