From 4f58318d5e566203158347961e7df598fe22907f Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 19 Jul 2026 03:17:59 +0530 Subject: [PATCH 1/6] use the Bijectors batched RQS for nsf and drop MonotonicSplines The neural spline coupling now calls Bijectors.rqs_params_from_raw, rqs_forward, and rqs_inverse, which are written with whole-array operations. nsf now trains under every supported AD backend and runs on the GPU, so the Zygote-only restriction is gone. Co-Authored-By: Claude Opus 4.8 (1M context) --- Project.toml | 8 +++--- README.md | 1 - docs/src/NSF.md | 5 ++-- example/demo_neural_spline_flow.jl | 5 ++-- src/NormalizingFlows.jl | 1 - src/flows/neuralspline.jl | 39 +++++++++++++----------------- test/Project.toml | 3 +-- test/ad.jl | 15 +++++------- 8 files changed, 33 insertions(+), 44 deletions(-) diff --git a/Project.toml b/Project.toml index 4c49b8f1..c123e5c7 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" +Bijectors = "0.16" CUDA = "5" 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/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 6e9455f9..6ab91bbe 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" @@ -18,4 +17,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Mooncake = "0.4.142" +Mooncake = "0.4.142, 0.5" diff --git a/test/ad.jl b/test/ad.jl index b8a42117..1e29c427 100644 --- a/test/ad.jl +++ b/test/ad.jl @@ -123,16 +123,13 @@ 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?) 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.AutoEnzyme(; + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), + function_annotation=Enzyme.Const, + ), + ADTypes.AutoMooncake(; config=Mooncake.Config()), ] @testset "$T" for T in [Float32, Float64] μ = 10 * ones(T, 2) From dc18aec3f76d46146d6fe367768ba025bfc6695d Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 19 Jul 2026 11:20:25 +0530 Subject: [PATCH 2/6] point the example environment at the Bijectors branch too The neural spline demo needs the batched RQS, which is not yet released, so the example environment needs the same source pin as the main project. Co-Authored-By: Claude Opus 4.8 (1M context) --- example/Project.toml | 4 ++++ 1 file changed, 4 insertions(+) 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" From 0c1bb27d0feb19b7d3872d31a4d8a947d4484386 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 20 Jul 2026 13:36:03 +0530 Subject: [PATCH 3/6] make the Julia 1.10 job install the Bijectors branch and cap Enzyme below 0.13.189 Julia 1.10 ignores the [sources] pin, so the batched RQS was missing there. Enzyme 0.13.189 fails with an LLVM canonicalization error on 1.10; 0.13.188 is fine. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/CI.yml | 3 +++ test/Project.toml | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) 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/test/Project.toml b/test/Project.toml index 827f8940..3978b5e4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,7 +21,7 @@ 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" From 66a23472376d7b92342fb53ef62fef9ad920ef3e Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 20 Jul 2026 14:08:32 +0530 Subject: [PATCH 4/6] clean up test warnings from repeated @leaf and a deprecated MvNormal constructor @leaf MvNormal now runs once in runtests.jl instead of inside every testset loop, and the base distributions use an identity covariance instead of the deprecated vector-of-sigmas form. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/ad.jl | 15 +++------------ test/flow.jl | 4 ---- test/interface.jl | 3 +-- test/objectives.jl | 2 +- test/runtests.jl | 3 +++ 5 files changed, 8 insertions(+), 19 deletions(-) diff --git a/test/ad.jl b/test/ad.jl index 1e29c427..a9c5ffec 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), I) 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), I) flow = realnvp(q₀, [8, 8], 3; paramtype=T) θ, re = Optimisers.destructure(flow) @@ -137,10 +131,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), I) 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..c12176e5 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), I) 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..3686f83c 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), I) 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") From aa5238dc7c49eef25acbec72933b1842444d4251 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 20 Jul 2026 14:08:32 +0530 Subject: [PATCH 5/6] run the NSF Enzyme AD test only on Julia 1.11 and newer Enzyme 0.13.188 fails LLVM function verification when differentiating the batched RQS on Julia 1.10 (Float64 only); the other backends pass there and Enzyme passes on 1.11+. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/ad.jl | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/test/ad.jl b/test/ad.jl index a9c5ffec..33d22496 100644 --- a/test/ad.jl +++ b/test/ad.jl @@ -116,15 +116,22 @@ end end @testset "AD for ELBO on NSF" begin - @testset "$at" for at in [ + nsf_adtypes = ADTypes.AbstractADType[ ADTypes.AutoZygote(), ADTypes.AutoReverseDiff(; compile=false), - ADTypes.AutoEnzyme(; - mode=Enzyme.set_runtime_activity(Enzyme.Reverse), - function_annotation=Enzyme.Const, - ), 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)) From e6684be60db6e8406e23df36fa878ab0eae09173 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 20 Jul 2026 14:23:02 +0530 Subject: [PATCH 6/6] keep the diagonal covariance type when replacing the deprecated MvNormal constructor MvNormal(mu, I) builds an IsoNormal, and Enzyme crashes differentiating destructure over that type. Diagonal(ones(T, 2)) keeps the exact DiagNormal the deprecated form produced. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/ad.jl | 6 +++--- test/interface.jl | 2 +- test/objectives.jl | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/test/ad.jl b/test/ad.jl index 33d22496..6ae70434 100644 --- a/test/ad.jl +++ b/test/ad.jl @@ -43,7 +43,7 @@ end target = MvNormal(μ, Σ) logp(z) = logpdf(target, z) - q₀ = MvNormal(zeros(T, 2), I) + q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2))) flow = Bijectors.transformed( q₀, Bijectors.Shift(zeros(T, 2)) ∘ Bijectors.Scale(ones(T, 2)) ) @@ -89,7 +89,7 @@ end target = MvNormal(μ, Σ) logp(z) = logpdf(target, z) - q₀ = MvNormal(zeros(T, 2), I) + q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2))) flow = realnvp(q₀, [8, 8], 3; paramtype=T) θ, re = Optimisers.destructure(flow) @@ -138,7 +138,7 @@ end target = MvNormal(μ, Σ) logp(z) = logpdf(target, z) - q₀ = MvNormal(zeros(T, 2), I) + 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/interface.jl b/test/interface.jl index c12176e5..5fbb7820 100644 --- a/test/interface.jl +++ b/test/interface.jl @@ -18,7 +18,7 @@ target = MvNormal(μ, Σ) logp(z) = logpdf(target, z) - q₀ = MvNormal(zeros(T, 2), I) + 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 3686f83c..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), I) + q₀ = MvNormal(zeros(T, 2), Diagonal(ones(T, 2))) flow = Bijectors.transformed(q₀, Bijectors.Shift(μ) ∘ Bijectors.Scale(sqrt.(Σ))) x = randn(T, 2)