make realnvp and nsf layers as part of the pkg #53
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Member
Author
|
After the recent Mooncake update ( using Random, Distributions, LinearAlgebra
using Bijectors
using Bijectors: partition, combine, PartitionMask
using Mooncake, Enzyme, ADTypes
import DifferentiationInterface as DI
# just define a MLP
function mlp3(
input_dim::Int,
hidden_dims::Int,
output_dim::Int;
activation=Flux.leakyrelu,
paramtype::Type{T} = Float64
) where {T<:AbstractFloat}
m = Chain(
Flux.Dense(input_dim, hidden_dims, activation),
Flux.Dense(hidden_dims, hidden_dims, activation),
Flux.Dense(hidden_dims, output_dim),
)
return Flux._paramtype(paramtype, m)
end
inputdim = 4
mask_idx = 1:2:inputdim
# creat a masking layer
mask = PartitionMask(inputdim, mask_idx)
cdim = length(mask_idx)
x = randn(inputdim)
t_net = mlp3(cdim, 16, cdim; paramtype = Float64)
ps, st = Optimisers.destructure(t_net)the following code runs perfectly function loss(ps, st, x, mask)
t_net = st(ps)
x₁, x₂, x₃ = partition(mask, x)
y₁ = x₁ .+ t_net(x₂)
y = combine(mask, y₁, x₂, x₃)
# println("y = ", y)
return sum(abs2, y)
end
loss(ps, st, x, mask) # return 3.0167880799441793
val, grad = DI.value_and_gradient(
ls_msk,
ADTypes.AutoMooncake(; config = Mooncake.Config()),
ps, DI.Cache(st), DI.Constant(x), DI.Constant(mask)
)but autograd fails if I wrap struct ACL
mask::Bijectors.PartitionMask
t::Flux.Chain
end
@functor ACL (t, )
acl = ACL(mask, t_net)
psacl, stacl = Optimisers.destructure(acl)
function loss_acl(ps, st, x)
acl = st(ps)
t_net = acl.t
mask = acl.mask
x₁, x₂, x₃ = partition(mask, x)
y₁ = x₁ .+ t_net(x₂)
y = combine(mask, y₁, x₂, x₃)
return sum(abs2, y)
end
loss_acl(psacl, stacl, x) # return 3.0167880799441793
val, grad = DI.value_and_gradient(
loss_acl,
ADTypes.AutoMooncake(; config = Mooncake.Config()),
psacl, DI.Cache(stacl), DI.Constant(x)
)with error message
val, grad = DI.value_and_gradient(
loss_acl,
ADTypes.AutoEnzyme(;
mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
function_annotation=Enzyme.Const,
),
psacl, DI.Cache(stacl), DI.Constant(x)
)with output Any thoughts on this @yebai @willtebbutt? |
Member
Author
|
Ah looks like it only has issue when the part of the fields in the structure is annotated by struct Holder
t::Flux.Chain
end
@functor Holder
psh, sth = Optimisers.destructure(Holder(t_net))
function loss2(ps, st, x)
holder = st(ps)
t_net = holder.t
y = x .+ t_net(x)
return sum(abs2, y)
end
loss2(psh, sth, x) # return 7.408352005690478
val, grad = DI.value_and_gradient(
loss2,
ADTypes.AutoMooncake(; config = Mooncake.Config()),
psh, DI.Cache(sth), DI.Constant(x)
)with outputs |
Member
|
@zuhengxu, can you help bisect which Mooncake version / Julia version introduced this bug? |
Member
Author
Good point! I'll look at this today. |
Member
Author
Member
|
It appears that the remaining issues with Mooncake are minor, likely due to a lack of a specific rule. @sunxd3, can you help if it requires a new rule? |
Member
|
I'll look into it 👍 |
Red-Portal
requested changes
Jul 31, 2025
Member
Red-Portal
left a comment
Red-Portal
left a comment
There was a problem hiding this comment.
Hi! I have some minor suggestions
Contributor
|
NormalizingFlows.jl documentation for PR #53 is available at: |
Member
Author
|
Thank you @yebai @sunxd3 @Red-Portal again for the help and comments in the process of this PR! Let me know if this PR looks good to you and I'll merge it afterwards. |
Member
|
Sorry for the delay. Reviewing a paper by JMLR has been taking up all my bandwidth. I'll take a deeper look tomorrow. |
Red-Portal
requested changes
Aug 14, 2025
Member
Red-Portal
left a comment
Red-Portal
left a comment
There was a problem hiding this comment.
I only have minor suggestions. Feel free to take a look and apply them if you agree. Otherwise, looks good to me.
Red-Portal
reviewed
Aug 14, 2025
Member
|
sorry for missing the tag, allow me to give a look later today or tomorrow |
sunxd3
reviewed
Aug 18, 2025
Member
sunxd3
left a comment
sunxd3
left a comment
There was a problem hiding this comment.
couple of tiny things, very happy to do another round of review
Member
Author
|
Thank you @sunxd3 @Red-Portal for the review! I made the corresponding updates and let me know if the current version looks good to you! |
sunxd3
reviewed
Aug 20, 2025
Member
sunxd3
left a comment
sunxd3
left a comment
There was a problem hiding this comment.
another couple of tiny things, nothing major beyond these
Member
|
pretty much good to go from my end, but let's wait for Kyurae to take a look? |
Red-Portal
requested changes
Aug 20, 2025
Member
Red-Portal
left a comment
Red-Portal
left a comment
There was a problem hiding this comment.
I only have a few minor comments that should be quick to handle!
Member
Author
|
@Red-Portal @sunxd3 Let me know if I can hit the big green button! Thanks for the quick feedback. |
Member
|
Alright looks good to me now! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As discussed in #36 (see #36 (comment)), I'm moving the
AffineCouplingandNeuralSplineLayerfrom the example tosrc/so it can be called.AffineCouplingandNeuralSplineLayerintosrcrealnvpand aneuralsplineflowconstructor. For therealnvp, follow the default architecture as mentioned in Advances in Black-Box VI: Normalizing Flows, Importance Weighting, and Optimization.