Replace Zygote by Enzyme

Project.toml

@@ -31,7 +31,6 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 ScientificTypesBase = "30f210dd-8aff-4c5f-94ba-8e64358c1161"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -68,5 +67,4 @@ SciMLBase = "2"
 SciMLSensitivity = "7"
 ScientificTypesBase = "3"
 Statistics = "1"
-Zygote = "0.6, 0.7"
 julia = "1.10"

README.md

@@ -74,13 +74,13 @@ r = rand(data_dist, nvars, n)
 r = convert.(Float32, r)
 
 # Fit It
-using DataFrames, MLJBase #, Zygote, ADTypes, OptimizationOptimisers
+using DataFrames, MLJBase #, Enzyme, ADTypes, OptimizationOptimisers
 df = DataFrame(transpose(r), :auto)
 model = ICNFModel(
     icnf;
     # optimizers = (Lion(),),
     # n_epochs = 300,
-    # adtype = AutoZygote(),
+    # adtype = AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
     # use_batch = true,
     # batch_size = 32,
     sol_kwargs = (; progress = true,), # pass to the solver

benchmark/Project.toml

@@ -7,7 +7,6 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ADTypes = "1"
@@ -18,5 +17,4 @@ Enzyme = "0.13"
 Lux = "1"
 PkgBenchmark = "0.2"
 StableRNGs = "1"
-Zygote = "0.6, 0.7"
 julia = "1.10"

benchmark/benchmarks.jl

@@ -6,7 +6,6 @@ import ADTypes,
     Lux,
     PkgBenchmark,
     StableRNGs,
-    Zygote,
     ContinuousNormalizingFlows
 
 SUITE = BenchmarkTools.BenchmarkGroup()
@@ -58,8 +57,16 @@ end
 
 diff_loss_tn(ps)
 diff_loss_tt(ps)
-DifferentiationInterface.gradient(diff_loss_tn, ADTypes.AutoZygote(), ps)
-DifferentiationInterface.gradient(diff_loss_tt, ADTypes.AutoZygote(), ps)
+DifferentiationInterface.gradient(
+    diff_loss_tn,
+    ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
+    ps,
+)
+DifferentiationInterface.gradient(
+    diff_loss_tt,
+    ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
+    ps,
+)
 GC.gc()
 
 SUITE["main"]["no_inplace"]["direct"]["train"] =
@@ -69,13 +76,13 @@ SUITE["main"]["no_inplace"]["direct"]["test"] =
 SUITE["main"]["no_inplace"]["AD-1-order"]["train"] =
     BenchmarkTools.@benchmarkable DifferentiationInterface.gradient(
         diff_loss_tn,
-        ADTypes.AutoZygote(),
+        ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
         ps,
     )
 SUITE["main"]["no_inplace"]["AD-1-order"]["test"] =
     BenchmarkTools.@benchmarkable DifferentiationInterface.gradient(
         diff_loss_tt,
-        ADTypes.AutoZygote(),
+        ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
         ps,
     )
 
@@ -106,8 +113,16 @@ end
 
 diff_loss_tn2(ps)
 diff_loss_tt2(ps)
-DifferentiationInterface.gradient(diff_loss_tn2, ADTypes.AutoZygote(), ps)
-DifferentiationInterface.gradient(diff_loss_tt2, ADTypes.AutoZygote(), ps)
+DifferentiationInterface.gradient(
+    diff_loss_tn2,
+    ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
+    ps,
+)
+DifferentiationInterface.gradient(
+    diff_loss_tt2,
+    ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
+    ps,
+)
 GC.gc()
 
 SUITE["main"]["inplace"]["direct"]["train"] =
@@ -116,12 +131,12 @@ SUITE["main"]["inplace"]["direct"]["test"] = BenchmarkTools.@benchmarkable diff_
 SUITE["main"]["inplace"]["AD-1-order"]["train"] =
     BenchmarkTools.@benchmarkable DifferentiationInterface.gradient(
         diff_loss_tn2,
-        ADTypes.AutoZygote(),
+        ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
         ps,
     )
 SUITE["main"]["inplace"]["AD-1-order"]["test"] =
     BenchmarkTools.@benchmarkable DifferentiationInterface.gradient(
         diff_loss_tt2,
-        ADTypes.AutoZygote(),
+        ADTypes.AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)),
         ps,
     )

src/ContinuousNormalizingFlows.jl

@@ -27,8 +27,7 @@ import ADTypes,
     ScientificTypesBase,
     SciMLBase,
     SciMLSensitivity,
-    Statistics,
-    Zygote
+    Statistics
 
 export construct,
     inference,

src/exts/mlj_ext/core_cond_icnf.jl

@@ -16,7 +16,9 @@ function CondICNFModel(
     loss::Function = loss;
     optimizers::Tuple = (Optimisers.Lion(),),
     n_epochs::Int = 300,
-    adtype::ADTypes.AbstractADType = ADTypes.AutoZygote(),
+    adtype::ADTypes.AbstractADType = ADTypes.AutoEnzyme(;
+        mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+    ),
     use_batch::Bool = true,
     batch_size::Int = 32,
     sol_kwargs::NamedTuple = (;),

src/exts/mlj_ext/core_icnf.jl

@@ -16,7 +16,9 @@ function ICNFModel(
     loss::Function = loss;
     optimizers::Tuple = (Optimisers.Lion(),),
     n_epochs::Int = 300,
-    adtype::ADTypes.AbstractADType = ADTypes.AutoZygote(),
+    adtype::ADTypes.AbstractADType = ADTypes.AutoEnzyme(;
+        mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+    ),
     use_batch::Bool = true,
     batch_size::Int = 32,
     sol_kwargs::NamedTuple = (;),

src/utils.jl

@@ -6,14 +6,14 @@
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = similar(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[i, :, :] =
             only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (z,)))
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
     end
-    y, eachslice(copy(res); dims = 3)
+    y, eachslice(res; dims = 3)
 end
 
 @inline function jacobian_batched(
@@ -24,15 +24,15 @@ end
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = similar(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[:, i, :] = only(
             DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (z,)),
         )
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
     end
-    y, eachslice(copy(res); dims = 3)
+    y, eachslice(res; dims = 3)
 end
 
 @inline function jacobian_batched(