Use the Bijectors batched RQS for nsf and drop MonotonicSplines

.github/workflows/CI.yml

@@ -34,5 +34,8 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v3
+      - name: Pin Bijectors to the branch on Julia 1.10
+        if: matrix.version == 'min'
+        run: julia --project=. --color=yes -e 'using Pkg; Pkg.add(PackageSpec(url="https://github.com/TuringLang/Bijectors.jl", rev="sg/batched-gpu-rqs"))'
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1

Project.toml

@@ -11,7 +11,6 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MonotonicSplines = "568f7cb4-8305-41bc-b90d-d32b39cc99d1"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -23,16 +22,19 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 [extensions]
 NormalizingFlowsCUDAExt = "CUDA"
 
+# Temporary: batched RQS is on an unreleased Bijectors branch. Remove once it is tagged.
+[sources]
+Bijectors = {url = "https://github.com/TuringLang/Bijectors.jl", rev = "sg/batched-gpu-rqs"}
+
 [compat]
 ADTypes = "1"
-Bijectors = "0.12.6, 0.13, 0.14, 0.15, 0.16"
+Bijectors = "0.16"
 CUDA = "5, 6.2"
 DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 Flux = "0.16"
 Functors = "0.5.2"
-MonotonicSplines = "0.3.3"
 Optimisers = "0.2.16, 0.3, 0.4"
 ProgressMeter = "1.0.0"
 StatsBase = "0.33, 0.34"

README.md

@@ -92,4 +92,3 @@ where one wants to learn the underlying distribution of some data.
 - [Flux.jl](https://fluxml.ai/Flux.jl/stable/)
 - [Optimisers.jl](https://github.com/FluxML/Optimisers.jl)
 - [AdvancedVI.jl](https://github.com/TuringLang/AdvancedVI.jl)
-- [MonotonicSplines.jl](https://github.com/bat/MonotonicSplines.jl)

docs/src/NSF.md

@@ -4,7 +4,7 @@
 using Random, Distributions, LinearAlgebra
 using Functors
 using Optimisers, ADTypes
-using Zygote
+using Mooncake
 using NormalizingFlows
 
 
@@ -27,8 +27,7 @@ sample_per_iter = 64
 
 # callback function to log training progress
 cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
-# nsf only supports AutoZygote
-adtype = ADTypes.AutoZygote()
+adtype = ADTypes.AutoMooncake(; config=Mooncake.Config())
 checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
     elbo_batch,

example/Project.toml

@@ -21,6 +21,10 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+# Temporary: batched RQS is on an unreleased Bijectors branch. Remove once it is tagged.
+[sources]
+Bijectors = {url = "https://github.com/TuringLang/Bijectors.jl", rev = "sg/batched-gpu-rqs"}
+
 [compat]
 ADTypes = "1.22.2"
 BenchmarkTools = "1.8.0"

example/demo_neural_spline_flow.jl

@@ -1,7 +1,7 @@
 using Random, Distributions, LinearAlgebra
 using Functors
 using Optimisers, ADTypes
-using Zygote
+using Mooncake
 using NormalizingFlows
 
 include("SyntheticTargets.jl")
@@ -36,8 +36,7 @@ sample_per_iter = 64
 
 # callback function to log training progress
 cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,ad=adtype)
-# nsf only supports AutoZygote
-adtype = ADTypes.AutoZygote()
+adtype = ADTypes.AutoMooncake(; config=Mooncake.Config())
 checkconv(iter, stat, re, θ, st) = stat.gradient_norm < one(T)/1000
 flow_trained, stats, _ = train_flow(
     elbo_batch,

src/NormalizingFlows.jl

@@ -132,7 +132,6 @@ include("flows/utils.jl")
 include("flows/planar_radial.jl")
 include("flows/realnvp.jl")
 
-using MonotonicSplines
 include("flows/neuralspline.jl")
 
 export create_flow

src/flows/neuralspline.jl

@@ -24,11 +24,9 @@ Fields
 
 Notes
 - Output dimensionality of the conditioner is `(3K - 1) * n_transformed`.
-- For computation performance, we rely on 
-[`MonotonicSplines.jl`](https://github.com/bat/MonotonicSplines.jl) for the
-building the rational quadratic spline functions.
-- See `MonotonicSplines.rqs_forward` and `MonotonicSplines.rqs_inverse` for forward/inverse 
-and log-determinant computations.
+- The rational quadratic spline is evaluated with the batched implementation in
+`Bijectors.jl`. See `Bijectors.rqs_forward` and `Bijectors.rqs_inverse` for the
+forward/inverse maps and log-determinant computations.
 
 [^DBMP2019]: Durkan, C., Bekasov, A., Murray, I. and Papamarkou, T. (2019). Neural Spline Flows. *NeurIPS.*
 """
@@ -64,10 +62,9 @@ end
 
 function get_nsc_params(nsc::NeuralSplineCoupling, x::AbstractVecOrMat)
     nnoutput = nsc.nn(x)
-    px, py, dydx = MonotonicSplines.rqs_params_from_nn(
-        nnoutput, nsc.n_dims_transformed, nsc.B
+    return Bijectors.rqs_params_from_raw(
+        _ensure_matrix(nnoutput), nsc.n_dims_transformed, nsc.B
     )
-    return px, py, dydx
 end
 
 # when input x is a vector instead of a matrix
@@ -80,14 +77,14 @@ function Bijectors.transform(nsc::NeuralSplineCoupling, x::AbstractVector)
     # 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)
+    y1, _ = Bijectors.rqs_forward(x1, px, py, dydx)
     return Bijectors.combine(nsc.mask, vec(y1), x2, x3)
 end
 function Bijectors.transform(nsc::NeuralSplineCoupling, x::AbstractMatrix)
     x1, x2, x3 = Bijectors.partition(nsc.mask, x)
     # instantiate rqs knots and derivatives
     px, py, dydx = get_nsc_params(nsc, x2)
-    y1, _ = MonotonicSplines.rqs_forward(x1, px, py, dydx)
+    y1, _ = Bijectors.rqs_forward(x1, px, py, dydx)
     return Bijectors.combine(nsc.mask, y1, x2, x3)
 end
 
@@ -96,14 +93,14 @@ function Bijectors.with_logabsdet_jacobian(nsc::NeuralSplineCoupling, x::Abstrac
     # 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)
+    y1, logjac = Bijectors.rqs_forward(x1, px, py, dydx)
     return Bijectors.combine(nsc.mask, vec(y1), x2, x3), logjac[1]
 end
 function Bijectors.with_logabsdet_jacobian(nsc::NeuralSplineCoupling, x::AbstractMatrix)
     x1, x2, x3 = Bijectors.partition(nsc.mask, x)
     # instantiate rqs knots and derivatives
     px, py, dydx = get_nsc_params(nsc, x2)
-    y1, logjac = MonotonicSplines.rqs_forward(x1, px, py, dydx)
+    y1, logjac = Bijectors.rqs_forward(x1, px, py, dydx)
     return Bijectors.combine(nsc.mask, y1, x2, x3), vec(logjac)
 end
 
@@ -112,14 +109,14 @@ function Bijectors.transform(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractV
     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)
+    x1, _ = Bijectors.rqs_inverse(y1, px, py, dydx)
     return Bijectors.combine(nsc.mask, vec(x1), y2, y3)
 end
 function Bijectors.transform(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractMatrix)
     nsc = insl.orig
     y1, y2, y3 = partition(nsc.mask, y)
     px, py, dydx = get_nsc_params(nsc, y2)
-    x1, _ = MonotonicSplines.rqs_inverse(y1, px, py, dydx)
+    x1, _ = Bijectors.rqs_inverse(y1, px, py, dydx)
     return Bijectors.combine(nsc.mask, x1, y2, y3)
 end
 
@@ -128,14 +125,14 @@ function Bijectors.with_logabsdet_jacobian(insl::Inverse{<:NeuralSplineCoupling}
     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)
+    x1, logjac = Bijectors.rqs_inverse(y1, px, py, dydx)
     return Bijectors.combine(nsc.mask, vec(x1), y2, y3), logjac[1]
 end
 function Bijectors.with_logabsdet_jacobian(insl::Inverse{<:NeuralSplineCoupling}, y::AbstractMatrix)
     nsc = insl.orig
     y1, y2, y3 = partition(nsc.mask, y)
     px, py, dydx = get_nsc_params(nsc, y2)
-    x1, logjac = MonotonicSplines.rqs_inverse(y1, px, py, dydx)
+    x1, logjac = Bijectors.rqs_inverse(y1, px, py, dydx)
     return Bijectors.combine(nsc.mask, x1, y2, y3), vec(logjac)
 end
 
@@ -204,12 +201,10 @@ Keyword Arguments
 Returns
 - `Bijectors.TransformedDistribution` representing the NSF flow.
 
-!!! note 
-    Under the hood, `nsf` relies on the rational quadratic spline function implememented in 
-    `MonotonicSplines.jl` for performance reasons.  `MonotonicSplines.jl` uses 
-    `KernelAbstractions.jl` to support batched operations. 
-    Because of this, so far `nsf` only supports `Zygote` as the AD type.
-
+!!! note
+    Under the hood, `nsf` uses the batched rational quadratic spline in `Bijectors.jl`. It is
+    written with whole-array operations, so the flow runs on the CPU and the GPU and is
+    differentiable by every supported AD backend.
 
 Example
 - `q0 = MvNormal(zeros(3), I); flow = nsf(q0, [64,64], 8, 3.0, 6)`

test/Project.toml

@@ -8,7 +8,6 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MonotonicSplines = "568f7cb4-8305-41bc-b90d-d32b39cc99d1"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -22,11 +21,10 @@ ADTypes = "1.22.2"
 Bijectors = "0.16.2"
 DifferentiationInterface = "0.7.20"
 Distributions = "0.25.129"
-Enzyme = "0.13.186"
+Enzyme = "0.13.186 - 0.13.188"
 Flux = "0.16.10"
 ForwardDiff = "1.4.1"
 Functors = "0.5.2"
-MonotonicSplines = "0.3.3"
 Mooncake = "0.4.142, 0.5"
 NormalizingFlows = "0.2.2"
 Optimisers = "0.4.7"

test/ad.jl

@@ -43,10 +43,7 @@ end
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
 
-            # necessary for Zygote/mooncake to differentiate through the flow
-            # prevent updating params of q0
-            @leaf MvNormal
-            q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+            q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2)))
             flow = Bijectors.transformed(
                 q₀, Bijectors.Shift(zeros(T, 2)) ∘ Bijectors.Scale(ones(T, 2))
             )
@@ -92,10 +89,7 @@ end
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
 
-            # necessary for Zygote/mooncake to differentiate through the flow
-            # prevent updating params of q0
-            @leaf MvNormal
-            q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+            q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2)))
             flow = realnvp(q₀, [8, 8], 3; paramtype=T)
 
             θ, re = Optimisers.destructure(flow)
@@ -122,28 +116,29 @@ end
 end
 
 @testset "AD for ELBO on NSF" begin
-    @testset "$at" for at in [
-        # now NSF only works with Zygote
-        # TODO: make it work with other ADs (possibly by adapting MonotonicSplines/src/rqspline_pullbacks.jl to rrules?)
+    nsf_adtypes = ADTypes.AbstractADType[
         ADTypes.AutoZygote(),
-        # ADTypes.AutoForwardDiff(),
-        # ADTypes.AutoReverseDiff(; compile=false),
-        # ADTypes.AutoEnzyme(;
-        #     mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
-        #     function_annotation=Enzyme.Const,
-        # ),
-        # ADTypes.AutoMooncake(; config=Mooncake.Config()),
+        ADTypes.AutoReverseDiff(; compile=false),
+        ADTypes.AutoMooncake(; config=Mooncake.Config()),
     ]
+    # Enzyme fails LLVM verification differentiating the batched RQS on Julia 1.10
+    if VERSION >= v"1.11"
+        push!(
+            nsf_adtypes,
+            ADTypes.AutoEnzyme(;
+                mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
+                function_annotation=Enzyme.Const,
+            ),
+        )
+    end
+    @testset "$at" for at in nsf_adtypes
         @testset "$T" for T in [Float32, Float64]
             μ = 10 * ones(T, 2)
             Σ = Diagonal(4 * ones(T, 2))
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
 
-            # necessary for Zygote/mooncake to differentiate through the flow
-            # prevent updating params of q0
-            @leaf MvNormal
-            q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+            q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2)))
             flow = nsf(q₀, [8, 8], 10, 5one(T), 3; paramtype=T)
 
             θ, re = Optimisers.destructure(flow)

test/flow.jl

@@ -7,7 +7,6 @@
     for T in [Float32, Float64]
         # Create a RealNVP flow
         q₀ = MvNormal(zeros(T, dim), I)
-        @leaf MvNormal
         flow = NormalizingFlows.realnvp(q₀, hdims, nlayers; paramtype=T)
 
         @testset "Sampling and density estimation for type: $T" begin
@@ -72,7 +71,6 @@ end
     for T in [Float32, Float64]
         # Create a nsf 
         q₀ = MvNormal(zeros(T, dim), I)
-        @leaf MvNormal
 
         B = 5one(T)
         flow = NormalizingFlows.nsf(q₀, hdims, K, B, nlayers; paramtype=T)
@@ -139,7 +137,6 @@ end
     for T in [Float32, Float64]
         # Create a nsf 
         q₀ = MvNormal(zeros(T, dim), I)
-        @leaf MvNormal
 
         flow = NormalizingFlows.planarflow(q₀, nlayers; paramtype=T)
 
@@ -205,7 +202,6 @@ end
     for T in [Float32, Float64]
         # Create a nsf 
         q₀ = MvNormal(zeros(T, dim), I)
-        @leaf MvNormal
 
         flow = NormalizingFlows.radialflow(q₀, nlayers; paramtype=T)
 

test/interface.jl

@@ -18,8 +18,7 @@
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
 
-            @leaf MvNormal
-            q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+            q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2)))
             flow = Bijectors.transformed(
                 q₀, Bijectors.Shift(zeros(T, 2)) ∘ Bijectors.Scale(ones(T, 2))
             )

test/objectives.jl

@@ -5,7 +5,7 @@
         target = MvNormal(μ, Σ)
         logp(z) = logpdf(target, z)
 
-        q₀ = MvNormal(zeros(T, 2), ones(T, 2))
+        q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2)))
         flow = Bijectors.transformed(q₀, Bijectors.Shift(μ) ∘ Bijectors.Scale(sqrt.(Σ)))
 
         x = randn(T, 2)

test/runtests.jl

@@ -11,6 +11,9 @@ import DifferentiationInterface as DI
 
 using Test
 
+# keep q0 parameters out of Optimisers.destructure
+@leaf MvNormal
+
 include("objectives.jl")
 include("interface.jl")
 include("flow.jl")