Apply suggestions from code review

Co-authored-by: David Widmann <[email protected]>

Author: sunxd3

src/NormalizingFlows.jl

@@ -11,7 +11,8 @@ using DocStringExtensions
 
 export train_flow, elbo, loglikelihood
 
-""" train_flow([rng::AbstractRNG, ]vo, flow, args...; kwargs...)
+""" 
+    train_flow([rng::AbstractRNG, ]vo, flow, args...; kwargs...)
 
 Train the given normalizing flow `flow` by calling `optimize`.
 
@@ -61,7 +62,7 @@ function train_flow(
         ADbackend,
         loss,
         θ_flat,
-        re, 
+        re,
         (rng, args...)...;
         max_iters=max_iters,
         optimiser=optimiser,

src/optimize.jl

@@ -2,23 +2,23 @@
 # training loop for variational objectives 
 #######################################################
 function pm_next!(pm, stats::NamedTuple)
-    return ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stats)])
+    return ProgressMeter.next!(pm; showvalues=map(tuple, keys(stats), values(stats)))
 end
 
-_wrap_in_DI_context(args) = DifferentiationInterface.Constant.([args...]) 
+_wrap_in_DI_context(args) = map(DifferentiationInterface.Constant, args)
 
 function _prepare_gradient(loss, adbackend, θ, args...)
     if isempty(args)
         return DifferentiationInterface.prepare_gradient(loss, adbackend, θ)
     end
-    return DifferentiationInterface.prepare_gradient(loss, adbackend, θ, _wrap_in_DI_context(args)...)
+    return DifferentiationInterface.prepare_gradient(loss, adbackend, θ, map(DifferentiationInterface.Constant, args)...)
 end
 
 function _value_and_gradient(loss, prep, adbackend, θ, args...)
     if isempty(args)
         return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ)
     end
-    return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ, _wrap_in_DI_context(args)...)
+    return DifferentiationInterface.value_and_gradient(loss, prep, adbackend, θ, map(DifferentiationInterface.Constant, args)...)
 end
 
 
@@ -42,7 +42,6 @@ Iteratively updating the parameters `θ` of the normalizing flow `re(θ)` by cal
 - `re`: reconstruction function that maps the flattened parameters to the normalizing flow
 - `args...`: additional arguments for `loss` (will be set as DifferentiationInterface.Constant)
 
-
 # Keyword Arguments
 - `max_iters::Int=10000`: maximum number of iterations
 - `optimiser::Optimisers.AbstractRule=Optimisers.ADAM()`: optimiser to compute the steps
@@ -102,9 +101,9 @@ function optimize(
             stat = (iteration=i, loss=ls, gradient_norm=norm(g))
 
             # callback
-            if !isnothing(callback)
+            if callback !== nothing
                 new_stat = callback(i, opt_stats, reconstruct, θ)
-                stat = !isnothing(new_stat) ? merge(stat, new_stat) : stat
+                stat = new_stat !== nothing ? merge(stat, new_stat) : stat
             end
             push!(opt_stats, stat)