dump the previous nsf implementation

Author: zuhengxu

.github/workflows/Examples.yml

@@ -38,7 +38,5 @@ jobs:
           include("demo_RealNVP.jl");
           @info "Running neural spline flow demo";
           include("demo_neural_spline_flow.jl");
-          @info "Running new neural spline flow demo";
-          include("demo_new_nsf.jl");
           @info "Running Hamiltonian flow demo";
           include("demo_hamiltonian_flow.jl");'

example/demo_neural_spline_flow.jl

@@ -4,7 +4,7 @@ using Bijectors: partition, combine, PartitionMask
 using Random, Distributions, LinearAlgebra
 using Functors
 using Optimisers, ADTypes
-using Mooncake
+using Zygote
 using NormalizingFlows
 
 include("SyntheticTargets.jl")
@@ -20,9 +20,9 @@ T = Float32
 ######################################
 # a difficult banana target
 ######################################
-
 target = Banana(2, one(T), 100one(T))
 logp = Base.Fix1(logpdf, target)
+
 ######################################
 # learn the target using Neural Spline Flow
 ######################################
@@ -39,14 +39,16 @@ sample_per_iter = 64
 
 # callback function to log training progress
 cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
-adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
+# TODO: mooncake has some issues with kernelabstractions?
+# adtype = ADTypes.AutoMooncake(; config = Mooncake.Config())
+adtype = ADTypes.AutoZygote()
 checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
-    elbo,
+    elbo_batch,
     flow,
     logp,
     sample_per_iter;
-    max_iters=10,   # change to larger number of iterations (e.g., 50_000) for better results
+    max_iters=10000,   # change to larger number of iterations (e.g., 50_000) for better results
     optimiser=Optimisers.Adam(1e-4),
     ADbackend=adtype,
     show_progress=true,

example/demo_new_nsf.jl

@@ -130,15 +130,10 @@ end
 include("flows/utils.jl")
 include("flows/realnvp.jl")
 include("flows/neuralspline.jl")
-# a new implementation of Neural Spline Flow based on MonotonicSplines.jl
-# the construction of the RQS seems to be more efficient than the one in Bijectors.jl
-# and supports batched operations.
-include("flows/new_nsf.jl")
 
 export create_flow
 export AffineCoupling, RealNVP_layer, realnvp
 export NeuralSplineCoupling, NSF_layer, nsf
-export NSC, new_NSF_layer, new_nsf
 
 
 end

src/NormalizingFlows.jl

@@ -1,6 +1,10 @@
+using MonotonicSplines
+# a new implementation of Neural Spline Flow based on MonotonicSplines.jl
+# the construction of the RQS seems to be more efficient than the one in Bijectors.jl
+# and supports batched operations.
+
 """
 Neural Rational quadratic Spline layer 
-
 # References
 [1] Durkan, C., Bekasov, A., Murray, I., & Papamakarios, G., Neural Spline Flows, CoRR, arXiv:1906.04032 [stat.ML],  (2019). 
 """
@@ -18,9 +22,9 @@ function NeuralSplineCoupling(
     hdims::AbstractVector{T1},       # dimension of hidden units for s and t
     K::T1,                           # number of knots
     B::T2,                           # bound of the knots
-    mask_idx::AbstractVector{<:Int}, # index of dimensione that one wants to apply transformations on
-    paramtype::Type{T3},             # type of the parameters, e.g., Float64 or Float32
-) where {T1<:Int,T2<:Real,T3<:AbstractFloat}
+    mask_idx::AbstractVector{T1}, # index of dimensione that one wants to apply transformations on
+    paramtype::Type{T2},             # type of the parameters, e.g., Float64 or Float32
+) where {T1<:Int,T2<:AbstractFloat}
     num_of_transformed_dims = length(mask_idx)
     input_dims = dim - num_of_transformed_dims
 
@@ -30,86 +34,75 @@ function NeuralSplineCoupling(
     nn = fnn(input_dims, hdims, output_dims; output_activation=nothing, paramtype=paramtype)
 
     mask = Bijectors.PartitionMask(dim, mask_idx)
-    return NeuralSplineCoupling(dim, K, num_of_transformed_dims, B, nn, mask)
+    return NeuralSplineCoupling{T2, typeof(nn)}(dim, K, num_of_transformed_dims, B, nn, mask)
 end
 
 @functor NeuralSplineCoupling (nn,)
 
-"""
-Build a rational quadratic spline (RQS) from the nn output
-Bijectors.jl has implemented the inverse and logabsdetjac for rational quadratic spline
-
-we just need to map the nn output to the knots and derivatives of the RQS
-"""
-function instantiate_rqs(nsl::NeuralSplineCoupling, x::AbstractVector)
-    K, B = nsl.K, nsl.B
-    nnoutput = reshape(nsl.nn(x), nsl.n_dims_transferred, :)
-    ws = @view nnoutput[:, 1:K]
-    hs = @view nnoutput[:, (K + 1):(2K)]
-    ds = @view nnoutput[:, (2K + 1):(3K - 1)]
-    return Bijectors.RationalQuadraticSpline(ws, hs, ds, B)
+function get_nsc_params(nsc::NeuralSplineCoupling, x::AbstractVecOrMat)
+    nnoutput = nsc.nn(x)
+    px, py, dydx = MonotonicSplines.rqs_params_from_nn(nnoutput, nsc.n_dims_transferred, nsc.B)
+    return px, py, dydx
 end
 
-function Bijectors.transform(nsl::NeuralSplineCoupling, x::AbstractVector)
-    x_1, x_2, x_3 = Bijectors.partition(nsl.mask, x)
-    # instantiate rqs knots and derivatives
-    rqs = instantiate_rqs(nsl, x_2)
-    y_1 = Bijectors.transform(rqs, x_1)
-    return Bijectors.combine(nsl.mask, y_1, x_2, x_3)
-end
+# when input x is a vector instead of a matrix
+# need this to transform it to a matrix with one row
+# otherwise, rqs_forward and rqs_inverse will throw an error
+_ensure_matrix(x) = x isa AbstractVector ? reshape(x, 1, length(x)) : x
 
-function Bijectors.transform(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractVector)
-    nsl = insl.orig
-    y1, y2, y3 = partition(nsl.mask, y)
-    rqs = instantiate_rqs(nsl, y2)
-    x1 = Bijectors.transform(Inverse(rqs), y1)
-    return Bijectors.combine(nsl.mask, x1, y2, y3)
+function Bijectors.transform(nsc::NeuralSplineCoupling, x::AbstractVecOrMat)
+    x1, x2, x3 = Bijectors.partition(nsc.mask, x)
+    # instantiate rqs knots and derivatives
+    px, py, dydx = get_nsc_params(nsc, x2)
+    x1 = _ensure_matrix(x1)
+    y1, _ = MonotonicSplines.rqs_forward(x1, px, py, dydx)
+    return Bijectors.combine(nsc.mask, y1, x2, x3)
 end
 
-function (nsl::NeuralSplineCoupling)(x::AbstractVector)
-    return Bijectors.transform(nsl, x)
+function Bijectors.with_logabsdet_jacobian(nsc::NeuralSplineCoupling, x::AbstractVecOrMat)
+    x1, x2, x3 = Bijectors.partition(nsc.mask, x)
+    # instantiate rqs knots and derivatives
+    px, py, dydx = get_nsc_params(nsc, x2)
+    x1 = _ensure_matrix(x1)
+    y1, logjac = MonotonicSplines.rqs_forward(x1, px, py, dydx)
+    return Bijectors.combine(nsc.mask, y1, x2, x3), logjac isa Real ? logjac : vec(logjac)
 end
 
-# define logabsdetjac
-function Bijectors.logabsdetjac(nsl::NeuralSplineCoupling, x::AbstractVector)
-    x_1, x_2, _ = Bijectors.partition(nsl.mask, x)
-    rqs = instantiate_rqs(nsl, x_2)
-    logjac = logabsdetjac(rqs, x_1)
-    return logjac
+function Bijectors.transform(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractVecOrMat)
+    nsc = insl.orig
+    y1, y2, y3 = partition(nsc.mask, y)
+    px, py, dydx = get_nsc_params(nsc, y2)
+    y1 = _ensure_matrix(y1)
+    x1, _ = MonotonicSplines.rqs_inverse(y1, px, py, dydx)
+    return Bijectors.combine(nsc.mask, x1, y2, y3)
 end
 
-function Bijectors.logabsdetjac(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractVector)
-    nsl = insl.orig
-    y1, y2, _ = partition(nsl.mask, y)
-    rqs = instantiate_rqs(nsl, y2)
-    logjac = logabsdetjac(Inverse(rqs), y1)
-    return logjac
+function Bijectors.with_logabsdet_jacobian(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractVecOrMat)
+    nsc = insl.orig
+    y1, y2, y3 = partition(nsc.mask, y)
+    px, py, dydx = get_nsc_params(nsc, y2)
+    y1 = _ensure_matrix(y1)
+    x1, logjac = MonotonicSplines.rqs_inverse(y1, px, py, dydx)
+    return Bijectors.combine(nsc.mask, x1, y2, y3), logjac isa Real ? logjac : vec(logjac)
 end
 
-function Bijectors.with_logabsdet_jacobian(nsl::NeuralSplineCoupling, x::AbstractVector)
-    x_1, x_2, x_3 = Bijectors.partition(nsl.mask, x)
-    rqs = instantiate_rqs(nsl, x_2)
-    y_1, logjac = with_logabsdet_jacobian(rqs, x_1)
-    return Bijectors.combine(nsl.mask, y_1, x_2, x_3), logjac
+function (nsc::NeuralSplineCoupling)(x::AbstractVecOrMat)
+    return Bijectors.transform(nsc, x)
 end
 
 
 """
     NSF_layer(dims, hdims; paramtype = Float64)
-
 Default constructor of single layer of Neural Spline Flow (NSF) 
 which is a composition of 2 neural spline coupling transformations with complementary masks.
 The masking strategy is odd-even masking.
-
 # Arguments
 - `dims::Int`: dimension of the problem
 - `hdims::AbstractVector{Int}`: dimension of hidden units for s and t
 - `K::Int`: number of knots
 - `B::AbstractFloat`: bound of the knots
-
 # Keyword Arguments
 - `paramtype::Type{T} = Float64`: type of the parameters, defaults to `Float64`
-
 # Returns
 - A `Bijectors.Bijector` representing the NSF layer.
 """