From 1b36c6e0ffbbe6b5bdf889f3fdcb29ca51e70e26 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Wed, 7 Aug 2024 11:14:10 +0200
Subject: [PATCH] Fix Enzyme extension (#79)

* Fix Enzyme extension

* Enable Enzyme tests

* Fix format

* Fix format

* Do not test on Julia nightly
---
 .github/workflows/JuliaNightly.yml | 33 ------------------------------
 README.md                          |  1 -
 ext/AdvancedVIEnzymeExt.jl         | 16 ++++++++++-----
 test/Project.toml                  |  2 ++
 test/interface/ad.jl               | 10 ++++-----
 test/runtests.jl                   |  2 +-
 6 files changed, 19 insertions(+), 45 deletions(-)
 delete mode 100644 .github/workflows/JuliaNightly.yml

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