Merge pull request #38 from ptiede/ptiede-enzymefix

Fix Enzyme implementation for 0.13

Author: tpapp

.github/workflows/CI.yml

@@ -17,15 +17,15 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.9' # Replace this with the minimum Julia version that your package supports. 
+          - 'min' # Minimum Julia version that the package supports. 
           - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
         os:
           - ubuntu-latest
         arch:
           - x64
     steps:
       - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@v1
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}

Project.toml

@@ -31,14 +31,14 @@ LogDensityProblemsADZygoteExt = "Zygote"
 ADTypes = "1.5"
 BenchmarkTools = "1"
 DocStringExtensions = "0.8, 0.9"
-Enzyme = "0.11, 0.12, 0.13"
+Enzyme = "0.13.3"
 FiniteDifferences = "0.12"
 ForwardDiff = "0.10"
 LogDensityProblems = "1, 2"
 ReverseDiff = "1"
 Tracker = "0.2"
 Zygote = "0.6"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/LogDensityProblemsADADTypesExt.jl

@@ -24,8 +24,13 @@ The AD configuration specified by `ad` is forwarded to the corresponding calls o
 """
 LogDensityProblemsAD.ADgradient(::ADTypes.AbstractADType, ℓ)
 
-function LogDensityProblemsAD.ADgradient(::ADTypes.AutoEnzyme, ℓ)
-    return LogDensityProblemsAD.ADgradient(Val(:Enzyme), ℓ)
+function LogDensityProblemsAD.ADgradient(ad::ADTypes.AutoEnzyme, ℓ)
+    if ad.mode === nothing
+        # Use default mode (Enzyme.Reverse)
+        return LogDensityProblemsAD.ADgradient(Val(:Enzyme), ℓ)
+    else
+        return LogDensityProblemsAD.ADgradient(Val(:Enzyme), ℓ; mode=ad.mode)
+    end
 end
 
 function LogDensityProblemsAD.ADgradient(ad::ADTypes.AutoForwardDiff{C}, ℓ) where {C}

ext/LogDensityProblemsADEnzymeExt.jl

@@ -8,6 +8,7 @@ if isdefined(Base, :get_extension)
 
     import LogDensityProblemsAD: ADgradient, logdensity_and_gradient
     import Enzyme
+    using Enzyme: EnzymeCore
 else
     using ..LogDensityProblemsAD: ADGradientWrapper, logdensity
 
@@ -46,7 +47,7 @@ function ADgradient(::Val{:Enzyme}, ℓ; mode::Enzyme.Mode = Enzyme.Reverse, sha
         @info "keyword argument `shadow` is ignored in reverse mode"
         shadow = nothing
     end
-    return EnzymeGradientLogDensity(ℓ, mode, shadow)
+    return EnzymeGradientLogDensity(ℓ, EnzymeCore.WithPrimal(mode), shadow)
 end
 
 function Base.show(io::IO, ∇ℓ::EnzymeGradientLogDensity)
@@ -58,17 +59,17 @@ function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme
                                  x::AbstractVector)
     (; ℓ, mode, shadow) = ∇ℓ
     _shadow = shadow === nothing ? Enzyme.onehot(x) : shadow
-    y, ∂ℓ_∂x = Enzyme.autodiff(mode, logdensity, Enzyme.BatchDuplicated,
+    ∂ℓ_∂x, y = Enzyme.autodiff(mode, logdensity, Enzyme.BatchDuplicated,
                                Enzyme.Const(ℓ),
                                Enzyme.BatchDuplicated(x, _shadow))
     return y, collect(∂ℓ_∂x)
 end
 
 function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme.ReverseMode},
                                  x::AbstractVector)
-    (; ℓ) = ∇ℓ
+    (; ℓ, mode) = ∇ℓ
     ∂ℓ_∂x = zero(x)
-    _, y = Enzyme.autodiff(Enzyme.ReverseWithPrimal, logdensity, Enzyme.Active,
+    _, y = Enzyme.autodiff(mode, logdensity, Enzyme.Active,
                            Enzyme.Const(ℓ), Enzyme.Duplicated(x, ∂ℓ_∂x))
     y, ∂ℓ_∂x
 end

test/runtests.jl

@@ -7,7 +7,7 @@ import ADTypes # load support for AD types with options
 import BenchmarkTools                            # load the heuristic chunks code
 using ComponentArrays: ComponentVector           # test with other vector types
 
-struct EnzymeTestMode <: Enzyme.Mode{Enzyme.DefaultABI, false} end
+struct EnzymeTestMode <: Enzyme.Mode{Enzyme.DefaultABI, false, false} end
 
 ####
 #### test setup and utilities
@@ -233,13 +233,18 @@ end
     ℓ = TestLogDensity(test_logdensity1)
 
     ∇ℓ_reverse = ADgradient(:Enzyme, ℓ)
+    ∇ℓ_forward = ADgradient(:Enzyme, ℓ; mode=Enzyme.Forward)
     @test ∇ℓ_reverse === ADgradient(:Enzyme, ℓ; mode=Enzyme.Reverse)
+    @test ∇ℓ_reverse.mode === Enzyme.ReverseWithPrimal
+    @test ADgradient(:Enzyme, ℓ; mode=Enzyme.Reverse) === ADgradient(:Enzyme, ℓ; mode=Enzyme.ReverseWithPrimal)
     @test repr(∇ℓ_reverse) == "Enzyme AD wrapper for " * repr(ℓ) * " with reverse mode"
 
     # ADTypes support
     @test ADgradient(ADTypes.AutoEnzyme(), ℓ) === ∇ℓ_reverse
 
     ∇ℓ_forward = ADgradient(:Enzyme, ℓ; mode=Enzyme.Forward)
+    @test ADgradient(ADTypes.AutoEnzyme(;mode=Enzyme.Forward), ℓ) === ∇ℓ_forward
+    @test ADgradient(:Enzyme, ℓ; mode=Enzyme.ForwardWithPrimal) === ADgradient(:Enzyme, ℓ; mode=Enzyme.ForwardWithPrimal)
     ∇ℓ_forward_shadow = ADgradient(:Enzyme, ℓ;
                                    mode=Enzyme.Forward,
                                    shadow=Enzyme.onehot(Vector{Float64}(undef, dimension(ℓ))))