fixing extension loading error

Author: zuhengxu

Project.toml

@@ -16,10 +16,10 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 
 [extensions]
-NormalizingFlowsEnzymeExt = "Enzyme"
+NormalizingFlowsEnzymeCoreExt = "EnzymeCore"
 
 [compat]
 ADTypes = "1"
@@ -32,3 +32,6 @@ ProgressMeter = "1.0.0"
 Requires = "1"
 StatsBase = "0.33, 0.34"
 julia = "1.10"
+
+[extras]
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"

ext/NormalizingFlowsEnzymeCoreExt.jl

@@ -0,0 +1,39 @@
+module NormalizingFlowsEnzymeCoreExt
+
+using EnzymeCore
+using NormalizingFlows
+using NormalizingFlows: ADTypes, DifferentiationInterface
+
+# deal with Enzyme readonly error: see https://discourse.julialang.org/t/enzyme-autodiff-readonly-error-and-working-with-batches-of-data/123012
+function NormalizingFlows._prepare_gradient(loss, adbackend::ADTypes.AutoEnzyme, θ, args...)
+    if isempty(args)
+        return DifferentiationInterface.prepare_gradient(
+            EnzymeCore.Const(loss), adbackend, θ
+        )
+    end
+    return DifferentiationInterface.prepare_gradient(
+        EnzymeCore.Const(loss),
+        adbackend,
+        θ,
+        map(DifferentiationInterface.Constant, args)...,
+    )
+end
+
+function NormalizingFlows._value_and_gradient(
+    loss, prep, adbackend::ADTypes.AutoEnzyme, θ, args...
+)
+    if isempty(args)
+        return DifferentiationInterface.value_and_gradient(
+            EnzymeCore.Const(loss), prep, adbackend, θ
+        )
+    end
+    return DifferentiationInterface.value_and_gradient(
+        EnzymeCore.Const(loss),
+        prep,
+        adbackend,
+        θ,
+        map(DifferentiationInterface.Constant, args)...,
+    )
+end
+
+end

ext/NormalizingFlowsEnzymeExt.jl

@@ -79,19 +79,17 @@ include("optimize.jl")
 include("objectives.jl")
 
 
-# optional dependencies 
-if !isdefined(Base, :get_extension) # check whether :get_extension is defined in Base
+if !isdefined(Base, :get_extension)
     using Requires
 end
 
-# Question: should Exts be loaded here or in train.jl? 
-function __init__()
-    @static if !isdefined(Base, :get_extension)
-        @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" include(
-            "../ext/NormalizingFlowsEnzymeExt.jl"
+
+@static if !isdefined(Base, :get_extension)
+    function __init__()
+        @require EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" include(
+            joinpath(@__DIR__, "../ext/NormalizingFlowsEnzymeCoreExt.jl")
         )
     end
 end
 
-
 end

src/NormalizingFlows.jl

@@ -93,7 +93,6 @@ function optimize(
         converged = false
         i = 1
         while (i ≤ max_iters) && !converged
-            # ls, g = DifferentiationInterface.value_and_gradient!(loss, grad, prep, adbackend, θ)
             ls, g = _value_and_gradient(loss, prep, adbackend, θ, args...)
 
             # Save stats

src/optimize.jl

@@ -4,6 +4,7 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/Project.toml

@@ -15,7 +15,6 @@
             ADTypes.AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse)),
             ADTypes.AutoMooncake(; config=Mooncake.Config()),
         ]
-            # at = ADTypes.AutoMooncake(; config=Mooncake.Config())
             prep = NormalizingFlows._prepare_gradient(f, at, x, y, z)
             value, grad = NormalizingFlows._value_and_gradient(f, prep, at, x, y, z)
             @test value ≈ f(x, y, z)
@@ -39,11 +38,13 @@ end
             logp(z) = logpdf(target, z)
 
             # necessary for Zygote/mooncake to differentiate through the flow
-            # prevent opt q0
+            # prevent updating params of q0
             @leaf MvNormal 
             q₀ = MvNormal(zeros(T, 2), ones(T, 2))
-            flow = Bijectors.transformed(q₀, Bijectors.Shift(zero.(μ)))
-
+            flow = Bijectors.transformed(
+                q₀, Bijectors.Shift(zeros(T, 2)) ∘ Bijectors.Scale(ones(T, 2))
+            )
+    
             θ, re = Optimisers.destructure(flow)
 
             # check grad computation for elbo

test/ad.jl

@@ -1,4 +1,4 @@
-@testset "learining 2d Gaussian" begin
+@testset "testing mean-field Gaussian VI" begin
     chunksize = 4
     @testset "$adtype" for adtype in [
         ADTypes.AutoZygote(),
@@ -48,3 +48,61 @@
         end
     end
 end
+
+# function create_planar_flow(n_layers::Int, q₀, T)
+#     d = length(q₀)
+#     if T == Float32
+#         Ls = reduce(∘, [f32(PlanarLayer(d)) for _ in 1:n_layers])
+#     else
+#         Ls = reduce(∘, [PlanarLayer(d) for _ in 1:n_layers])
+#     end
+#     return Bijectors.transformed(q₀, Ls)
+# end
+
+# @testset "testing planar flow" begin
+#     chunksize = 4
+#     @testset "$adtype" for adtype in [
+#         ADTypes.AutoZygote(),
+#         ADTypes.AutoForwardDiff(; chunksize=chunksize),
+#         ADTypes.AutoForwardDiff(),
+#         ADTypes.AutoReverseDiff(),
+#         ADTypes.AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse)),
+#         # ADTypes.AutoMooncake(; config = Mooncake.Config()), # somehow Mooncake does not work with Float64
+#     ]
+#         @testset "$T" for T in [Float32, Float64]
+#             μ = 10 * ones(T, 2)
+#             Σ = Diagonal(4 * ones(T, 2))
+
+#             target = MvNormal(μ, Σ)
+#             logp(z) = logpdf(target, z)
+
+#             @leaf MvNormal
+#             q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+#             nlayers = 10
+#             flow = create_planar_flow(nlayers, q₀, T)
+
+#             sample_per_iter = 10
+#             cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
+#             checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
+#             flow_trained, stats, _, _ = train_flow(
+#                 elbo,
+#                 flow,
+#                 logp,
+#                 sample_per_iter;
+#                 max_iters=10_000,
+#                 optimiser=Optimisers.Adam(one(T)/100),
+#                 ADbackend=adtype,
+#                 show_progress=false,
+#                 callback=cb,
+#                 hasconverged=checkconv,
+#             )
+#             θ, re = Optimisers.destructure(flow_trained)
+
+#             el_untrained = elbo(flow, logp, 1000)
+#             el_trained = elbo(flow_trained, logp, 1000)
+
+#             @test el_trained > el_untrained
+#             @test el_trained > -1
+#         end
+#     end
+# end

test/interface.jl

@@ -6,7 +6,9 @@ using Random
 using ADTypes
 using Functors
 using ForwardDiff, Zygote, ReverseDiff, Enzyme, Mooncake
+using Flux: f32
 import DifferentiationInterface as DI
+
 using Test
 
 include("ad.jl")