rm time_elapsed from train_flow

Author: zuhengxu

src/NormalizingFlows.jl

@@ -28,7 +28,13 @@ Train the given normalizing flow `flow` by calling `optimize`.
 - `optimiser::Optimisers.AbstractRule=Optimisers.ADAM()`: optimiser to compute the steps
 - `ADbackend::ADTypes.AbstractADType=ADTypes.AutoZygote()`: 
     automatic differentiation backend, currently supports
-    `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, and `ADTypes.ReverseDiff()`. 
+    `ADTypes.AutoZygote()`, `ADTypes.ForwardDiff()`, `ADTypes.ReverseDiff()`, 
+    `ADTypes.AutoMooncake()` and
+    `ADTypes.AutoEnzyme(;
+        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+        function_annotation=Enzyme.Const,
+    )`.
+    If user wants to use `AutoEnzyme`, please make sure to include the `set_runtime_activity` and `function_annotation` as shown above.
 - `kwargs...`: additional keyword arguments for `optimize` (See [`optimize`](@ref) for details)
 
 # Returns
@@ -58,7 +64,7 @@ function train_flow(
     loss(θ, rng, args...) = -vo(rng, re(θ), args...)
 
     # Normalizing flow training loop 
-    θ_flat_trained, opt_stats, st, time_elapsed = optimize(
+    θ_flat_trained, opt_stats, st = optimize(
         ADbackend,
         loss,
         θ_flat,
@@ -71,7 +77,7 @@ function train_flow(
     )
 
     flow_trained = re(θ_flat_trained)
-    return flow_trained, opt_stats, st, time_elapsed
+    return flow_trained, opt_stats, st
 end
 
 include("optimize.jl")

src/optimize.jl

@@ -13,7 +13,6 @@ function _value_and_gradient(loss, prep, adbackend, θ, args...)
     return DI.value_and_gradient(loss, prep, adbackend, θ, map(DI.Constant, args)...)
 end
 
-
 """
     optimize(
         ad::ADTypes.AbstractADType, 
@@ -58,7 +57,7 @@ Iteratively updating the parameters `θ` of the normalizing flow `re(θ)` by cal
 function optimize(
     adbackend,
     loss,
-    θ₀::AbstractVector{<:Real}, 
+    θ₀::AbstractVector{<:Real},
     reconstruct,
     args...;
     max_iters::Int=10000,
@@ -70,42 +69,40 @@ function optimize(
         max_iters; desc="Training", barlen=31, showspeed=true, enabled=show_progress
     ),
 )
-    time_elapsed = @elapsed begin 
-        opt_stats = []
+    opt_stats = []
 
-        # prepare loss and autograd
-        θ = deepcopy(θ₀)
-        # grad = similar(θ)
-        prep = _prepare_gradient(loss, adbackend, θ₀, args...)
+    # prepare loss and autograd
+    θ = deepcopy(θ₀)
+    # grad = similar(θ)
+    prep = _prepare_gradient(loss, adbackend, θ₀, args...)
 
-        # initialise optimiser state
-        st = Optimisers.setup(optimiser, θ)
+    # initialise optimiser state
+    st = Optimisers.setup(optimiser, θ)
 
-        # general `hasconverged(...)` approach to allow early termination.
-        converged = false
-        i = 1
-        while (i ≤ max_iters) && !converged
-            ls, g = _value_and_gradient(loss, prep, adbackend, θ, args...)
+    # general `hasconverged(...)` approach to allow early termination.
+    converged = false
+    i = 1
+    while (i ≤ max_iters) && !converged
+        ls, g = _value_and_gradient(loss, prep, adbackend, θ, args...)
 
-            # Save stats
-            stat = (iteration=i, loss=ls, gradient_norm=norm(g))
+        # Save stats
+        stat = (iteration=i, loss=ls, gradient_norm=norm(g))
 
-            # callback
-            if callback !== nothing
-                new_stat = callback(i, opt_stats, reconstruct, θ)
-                stat = new_stat !== nothing ? merge(stat, new_stat) : stat
-            end
-            push!(opt_stats, stat)
+        # callback
+        if callback !== nothing
+            new_stat = callback(i, opt_stats, reconstruct, θ)
+            stat = new_stat !== nothing ? merge(stat, new_stat) : stat
+        end
+        push!(opt_stats, stat)
 
-            # update optimiser state and parameters
-            st, θ = Optimisers.update!(st, θ, g)
+        # update optimiser state and parameters
+        st, θ = Optimisers.update!(st, θ, g)
 
-            # check convergence
-            i += 1
-            converged = hasconverged(i, stat, reconstruct, θ, st)
-            pm_next!(prog, stat)
-        end
+        # check convergence
+        i += 1
+        converged = hasconverged(i, stat, reconstruct, θ, st)
+        pm_next!(prog, stat)
     end
     # return status of the optimiser for potential continuation of training
-    return θ, map(identity, opt_stats), st, time_elapsed
+    return θ, map(identity, opt_stats), st
 end

test/interface.jl

@@ -27,7 +27,7 @@
             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(
+            flow_trained, stats, _ = train_flow(
                 elbo,
                 flow,
                 logp,