Error in example

[herluc]

Hello there,

I'm completely new to Normalizing Flows and this package is the first thing I try. I tried to run the example but it fails at

flow_trained, stats, _ = train_flow(
    elbo,
    flow,
    logp,
    sample_per_iter;
    max_iters=200_00,
    optimiser=Optimisers.ADAM(),
    callback=cb,
    hasconverged=checkconv,
    ADbackend=AutoZygote(), # using Zygote as the AD backend
)

with

ERROR: UndefVarError: `j` not defined in `Zygote`
Suggestion: check for spelling errors or missing imports.

Is there anything special that I have to configure? Using AutoMooncake instead, it works.

All the best
Lucas

[zuhengxu]

Thank you for trying our pkg! May I ask which example (demo file) you run? And which version of julia, Zygote were used?

[herluc]

julia> VERSION
v"1.11.5"

julia> pkgversion(Zygote)
v"0.6.77"

I just copy pasted the code from the "Example" section in the documentation

[zuhengxu]

@herluc Thanks for the info. The example should be working now with NormalizingFlows.jl (v0.2.1). Can you update NormalizingFlow.jl and see if the following code snippet runs (pwd NormalizingFlows.jl/example/)?

using Bijectors, FunctionChains
using Functors
using LinearAlgebra
using Flux: f32

using NormalizingFlows
using ADTypes
using Optimisers
using Zygote

# loads the definition of Banana target
include("SyntheticTargets.jl")

target = Banana(2, 1.0f-1, 100.0f0)
logp = Base.Fix1(logpdf, target)


# create planar flow
function create_planar_flow(n_layers::Int, q₀)
    d = length(q₀)
    Ls = [f32(PlanarLayer(d)) for _ in 1:n_layers]
    ts = fchain(Ls)
    return transformed(q₀, ts)
end

# create a 20-layer planar flow
@leaf MvNormal # to prevent params in q₀ from being optimized
q₀ = MvNormal(zeros(Float32, 2), I)
flow = create_planar_flow(20, q₀)


# training the flow
sample_per_iter = 10
# callback function to track the number of samples used per iteration
cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,)
# defined stopping criteria when the gradient norm is less than 1e-3
checkconv(iter, stat, re, θ, st) = stat.gradient_norm < 1e-3
flow_trained, stats, _ = train_flow(
    elbo,
    flow,
    logp,
    sample_per_iter;
    max_iters=200_00,
    optimiser=Optimisers.ADAM(),
    callback=cb,
    hasconverged=checkconv,
    ADbackend=AutoZygote(), # using Zygote as the AD backend
)

It works on my end with julia v1.11.5 and ` Zygote 0.6.77'.

[zuhengxu]

@herluc could you confirm whether the bug is fixed with NF.jl v0.2.1? Thanks!

[zuhengxu]

@herluc if there is no further problems after updating the pkg, I will close this issue as resolved.