diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 5dea1f20..3202f4ff 100644 --- a/.github/workflows/CI.yml +++ b/.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 diff --git a/Project.toml b/Project.toml index 109b2122..48b19b4f 100644 --- a/Project.toml +++ b/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" diff --git a/README.md b/README.md index 5bf5651f..14debf20 100644 --- a/README.md +++ b/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) diff --git a/docs/src/NSF.md b/docs/src/NSF.md index df03a410..1b7a61a0 100644 --- a/docs/src/NSF.md +++ b/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, diff --git a/example/Project.toml b/example/Project.toml index 8b44b7b3..73b27543 100644 --- a/example/Project.toml +++ b/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" diff --git a/example/demo_neural_spline_flow.jl b/example/demo_neural_spline_flow.jl index 4e98c308..ec3ccf78 100644 --- a/example/demo_neural_spline_flow.jl +++ b/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, diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl index 70c46996..f4616551 100644 --- a/src/NormalizingFlows.jl +++ b/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 diff --git a/src/flows/neuralspline.jl b/src/flows/neuralspline.jl index 0d8eab0c..dd3b2a40 100644 --- a/src/flows/neuralspline.jl +++ b/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)` diff --git a/test/Project.toml b/test/Project.toml index b45ccaa8..3978b5e4 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/ad.jl b/test/ad.jl index b8a42117..6ae70434 100644 --- a/test/ad.jl +++ b/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) diff --git a/test/flow.jl b/test/flow.jl index 69685588..46cc3b8d 100644 --- a/test/flow.jl +++ b/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) diff --git a/test/interface.jl b/test/interface.jl index 947d4f37..5fbb7820 100644 --- a/test/interface.jl +++ b/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)) ) diff --git a/test/objectives.jl b/test/objectives.jl index fa977e83..c8104726 100644 --- a/test/objectives.jl +++ b/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) diff --git a/test/runtests.jl b/test/runtests.jl index 2808fb12..e256798d 100644 --- a/test/runtests.jl +++ b/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")