rm unnecessary EnzymeCoreExt

Author: zuhengxu

Project.toml

@@ -14,12 +14,6 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
-[weakdeps]
-EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
-
-[extensions]
-NormalizingFlowsEnzymeCoreExt = ["EnzymeCore", "ADTypes", "DifferentiationInterface"]
-
 [compat]
 ADTypes = "1"
 Bijectors = "0.12.6, 0.13, 0.14, 0.15"
@@ -30,6 +24,3 @@ Optimisers = "0.2.16, 0.3, 0.4"
 ProgressMeter = "1.0.0"
 StatsBase = "0.33, 0.34"
 julia = "1.10"
-
-[extras]
-EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"

ext/NormalizingFlowsEnzymeCoreExt.jl

@@ -5,7 +5,7 @@ using Optimisers
 using LinearAlgebra, Random, Distributions, StatsBase
 using ProgressMeter
 using ADTypes
-using DifferentiationInterface
+import DifferentiationInterface as DI
 
 using DocStringExtensions
 

src/NormalizingFlows.jl

@@ -5,14 +5,12 @@ function pm_next!(pm, stats::NamedTuple)
     return ProgressMeter.next!(pm; showvalues=map(tuple, keys(stats), values(stats)))
 end
 
-_wrap_in_DI_context(args) = map(DifferentiationInterface.Constant, args)
-
 function _prepare_gradient(loss, adbackend, θ, args...)
-    return DifferentiationInterface.prepare_gradient(loss, adbackend, θ, map(DifferentiationInterface.Constant, args)...)
+    return DI.prepare_gradient(loss, adbackend, θ, map(DI.Constant, args)...)
 end
 
 function _value_and_gradient(loss, prep, adbackend, θ, args...)
-    return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ, map(DifferentiationInterface.Constant, args)...)
+    return DI.value_and_gradient(loss, prep, adbackend, θ, map(DI.Constant, args)...)
 end
 
 
@@ -34,7 +32,7 @@ Iteratively updating the parameters `θ` of the normalizing flow `re(θ)` by cal
 - `loss`: a general loss function θ -> loss(θ, args...) returning a scalar loss value that will be minimised
 - `θ₀::AbstractVector{T}`: initial parameters for the loss function (in the context of normalizing flows, it will be the flattened flow parameters)
 - `re`: reconstruction function that maps the flattened parameters to the normalizing flow
-- `args...`: additional arguments for `loss` (will be set as DifferentiationInterface.Constant)
+- `args...`: additional arguments for `loss` (will be set as DI.Constant)
 
 # Keyword Arguments
 - `max_iters::Int=10000`: maximum number of iterations

src/optimize.jl

@@ -12,7 +12,10 @@
             ADTypes.AutoForwardDiff(; chunksize=chunksize),
             ADTypes.AutoForwardDiff(),
             ADTypes.AutoReverseDiff(; compile=false),
-            ADTypes.AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse)),
+            ADTypes.AutoEnzyme(;
+                mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+                function_annotation=Enzyme.Const,
+            ),
             ADTypes.AutoMooncake(; config=Mooncake.Config()),
         ]
             prep = NormalizingFlows._prepare_gradient(f, at, x, y, z)
@@ -27,33 +30,40 @@ end
     @testset "$at" for at in [
         ADTypes.AutoZygote(),
         ADTypes.AutoForwardDiff(),
-        ADTypes.AutoReverseDiff(; compile = false),
-        ADTypes.AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse)),
+        ADTypes.AutoReverseDiff(; compile=false),
+        ADTypes.AutoEnzyme(;
+            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+            function_annotation=Enzyme.Const,
+        ),
         ADTypes.AutoMooncake(; config=Mooncake.Config()),
     ]
         @testset "$T" for T in [Float32, Float64]
             μ = 10 * ones(T, 2)
             Σ = Diagonal(4 * ones(T, 2))
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
-            
+
             # necessary for Zygote/mooncake to differentiate through the flow
             # prevent updating params of q0
-            @leaf MvNormal 
+            @leaf MvNormal
             q₀ = MvNormal(zeros(T, 2), ones(T, 2))
             flow = Bijectors.transformed(
                 q₀, Bijectors.Shift(zeros(T, 2)) ∘ Bijectors.Scale(ones(T, 2))
             )
-    
+
             θ, re = Optimisers.destructure(flow)
 
             # check grad computation for elbo
-            loss(θ, rng, logp, sample_per_iter) = -NormalizingFlows.elbo(rng, re(θ), logp, sample_per_iter)
+            function loss(θ, rng, logp, sample_per_iter)
+                return -NormalizingFlows.elbo(rng, re(θ), logp, sample_per_iter)
+            end
 
             rng = Random.default_rng()
             sample_per_iter = 10
 
-            prep = NormalizingFlows._prepare_gradient(loss, at, θ, rng, logp, sample_per_iter)
+            prep = NormalizingFlows._prepare_gradient(
+                loss, at, θ, rng, logp, sample_per_iter
+            )
             value, grad = NormalizingFlows._value_and_gradient(
                 loss, prep, at, θ, rng, logp, sample_per_iter
             )

test/ad.jl

@@ -5,7 +5,10 @@
         ADTypes.AutoForwardDiff(; chunksize=chunksize),
         ADTypes.AutoForwardDiff(),
         ADTypes.AutoReverseDiff(),
-        ADTypes.AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse)),
+        ADTypes.AutoEnzyme(;
+            mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+            function_annotation=Enzyme.Const,
+        ),
         ADTypes.AutoMooncake(; config = Mooncake.Config()), 
     ]
         @testset "$T" for T in [Float32, Float64]