From 789d9e0459d65ad811b7f77002cf40360eeba1cd Mon Sep 17 00:00:00 2001
From: Zuheng <zuhengdavid@gmail.com>
Date: Fri, 20 Jun 2025 15:24:07 -0700
Subject: [PATCH 01/17] start testing full gpu training for realnvp

---
 example/Project.toml |   1 +
 example/test_gpu.jl  | 126 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 127 insertions(+)
 create mode 100644 example/test_gpu.jl

diff --git a/example/Project.toml b/example/Project.toml
index 0b9b0214..4ff5801d 100644
--- a/example/Project.toml
+++ b/example/Project.toml
@@ -2,6 +2,7 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
diff --git a/example/test_gpu.jl b/example/test_gpu.jl
new file mode 100644
index 00000000..f6da860d
--- /dev/null
+++ b/example/test_gpu.jl
@@ -0,0 +1,126 @@
+using Flux
+using Bijectors
+using Bijectors: partition, combine, PartitionMask
+
+using Random, Distributions, LinearAlgebra
+using Functors
+using Optimisers, ADTypes
+using Mooncake
+using NormalizingFlows
+
+include("SyntheticTargets.jl")
+include("utils.jl")
+
+##################################
+# define affine coupling layer using Bijectors.jl interface
+#################################
+struct AffineCoupling <: Bijectors.Bijector
+    dim::Int
+    mask::Bijectors.PartitionMask
+    s::Flux.Chain
+    t::Flux.Chain
+end
+
+# let params track field s and t
+@functor AffineCoupling (s, t)
+
+function AffineCoupling(
+    dim::Int,  # dimension of input
+    hdims::Int, # dimension of hidden units for s and t
+    mask_idx::AbstractVector, # index of dimensione that one wants to apply transformations on
+)
+    cdims = length(mask_idx) # dimension of parts used to construct coupling law
+    s = mlp3(cdims, hdims, cdims)
+    t = mlp3(cdims, hdims, cdims)
+    mask = PartitionMask(dim, mask_idx)
+    return AffineCoupling(dim, mask, s, t)
+end
+
+function Bijectors.transform(af::AffineCoupling, x::AbstractVecOrMat)
+    # partition vector using 'af.mask::PartitionMask`
+    x₁, x₂, x₃ = partition(af.mask, x)
+    y₁ = x₁ .* af.s(x₂) .+ af.t(x₂)
+    return combine(af.mask, y₁, x₂, x₃)
+end
+
+function (af::AffineCoupling)(x::AbstractArray)
+    return transform(af, x)
+end
+
+function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractVector)
+    x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
+    y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
+    logjac = sum(log ∘ abs, af.s(x_2)) # this is a scalar
+    return combine(af.mask, y_1, x_2, x_3), logjac
+end
+
+function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractMatrix)
+    x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
+    y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
+    logjac = sum(log ∘ abs, af.s(x_2); dims = 1) # 1 × size(x, 2)
+    return combine(af.mask, y_1, x_2, x_3), vec(logjac)
+end
+
+
+function Bijectors.with_logabsdet_jacobian(
+    iaf::Inverse{<:AffineCoupling}, y::AbstractVector
+)
+    af = iaf.orig
+    # partition vector using `af.mask::PartitionMask`
+    y_1, y_2, y_3 = partition(af.mask, y)
+    # inverse transformation
+    x_1 = (y_1 .- af.t(y_2)) ./ af.s(y_2)
+    logjac = -sum(log ∘ abs, af.s(y_2))
+    return combine(af.mask, x_1, y_2, y_3), logjac
+end
+
+function Bijectors.with_logabsdet_jacobian(
+    iaf::Inverse{<:AffineCoupling}, y::AbstractMatrix
+)
+    af = iaf.orig
+    # partition vector using `af.mask::PartitionMask`
+    y_1, y_2, y_3 = partition(af.mask, y)
+    # inverse transformation
+    x_1 = (y_1 .- af.t(y_2)) ./ af.s(y_2)
+    logjac = -sum(log ∘ abs, af.s(y_2); dims = 1)
+    return combine(af.mask, x_1, y_2, y_3), vec(logjac)
+end
+
+
+##################################
+# start demo
+#################################
+using CUDA
+const NF = NormalizingFlows
+rng_g = CUDA.default_rng() # use GPU RNG if available
+
+
+CUDA.allowscalar(true)
+n_samples = 100
+q0 = MvNormal(CUDA.zeros(2), cu([1f0 0f0; 0f0 1f0]))
+# gpu sample from the reference
+xs = NF._device_specific_rand(rng_g, q0, n_samples)
+
+d = 2
+hdims = 32
+Ls_g = [AffineCoupling(d, hdims, [1]) ∘ AffineCoupling(d, hdims, [2]) for i in 1:3]
+flow_g = create_flow(Ls_g, q0)
+flow_g = fmap(cu, flow_g) # move all flow parameters be on GPU
+
+# gpu sample from the flow
+ys = NF._device_specific_rand(rng_g, flow_g, n_samples)
+
+# log density computation
+logpdf(flow_g, ys) # errored
+
+logpdf(q0, xs) # returns a CPU array
+
+
+# elbo_batch(rng_g, flow, logp, n_samples)
+
+target = Banana(2, 1.0f0, 100.0f0)
+target_g = fmap(cu, target) # move target to GPU
+logp_g = Base.Fix1(logpdf, target_g)
+
+logp_g(yy)
+

From 98652be79f9362a755ab6d5964d8343ad35f06f4 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 02/17] Take the base distribution log-density on the device in
 the batched ELBO

---
 Project.toml                   |  6 +++++
 ext/NormalizingFlowsCUDAExt.jl | 16 ++++++++++----
 src/NormalizingFlows.jl        | 40 ++++++++++++++++++++++++++++++++++
 src/objectives/elbo.jl         |  2 +-
 4 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/Project.toml b/Project.toml
index b7528969..a8836a82 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,13 +6,16 @@ version = "0.4.0"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
@@ -30,12 +33,15 @@ ADTypes = "1.5"
 AbstractPPL = "0.15.4"
 Bijectors = "0.14.2, 0.15, 0.16"
 CUDA = "5, 6.2"
+ChainRulesCore = "1"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 Flux = "0.16"
 Functors = "0.5.2"
+GPUArraysCore = "0.1, 0.2"
 LogExpFunctions = "0.3.3"
 Optimisers = "0.2.16, 0.3, 0.4"
+PDMats = "0.11"
 ProgressMeter = "1.0.0"
 ReverseDiff = "1"
 StatsBase = "0.33, 0.34"
diff --git a/ext/NormalizingFlowsCUDAExt.jl b/ext/NormalizingFlowsCUDAExt.jl
index 1459cad7..b7d7e0c0 100644
--- a/ext/NormalizingFlowsCUDAExt.jl
+++ b/ext/NormalizingFlowsCUDAExt.jl
@@ -3,6 +3,7 @@ module NormalizingFlowsCUDAExt
 using CUDA
 using NormalizingFlows
 using NormalizingFlows: Bijectors, Distributions, Random
+using ChainRulesCore: @non_differentiable
 
 function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG,
@@ -23,9 +24,7 @@ function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
-    return @inbounds Distributions.rand!(
-        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, size(s))
-    )
+    return _cuda_draw(rng, s, size(s))
 end
 
 function _cuda_rand(
@@ -33,11 +32,20 @@ function _cuda_rand(
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
 )
+    return _cuda_draw(rng, s, (size(s)..., n))
+end
+
+function _cuda_draw(rng::CUDA.RNG, s, dims::Tuple)
     return @inbounds Distributions.rand!(
-        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, size(s)..., n)
+        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, dims)
     )
 end
 
+# Zygote cannot trace a `CuArray` allocation: it descends into the CUDA allocator and fails
+# to compile. Nothing is lost by hiding the draw, since `rand!` carries no gradient on the
+# host path either. Only the draw is opaque; a flow's transform stays differentiable.
+@non_differentiable _cuda_draw(::Any, ::Any, ::Any)
+
 # ! this is type piracy
 # replacing original function with scalar indexing
 function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVecOrMat)
diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl
index e24b5ffa..59043685 100644
--- a/src/NormalizingFlows.jl
+++ b/src/NormalizingFlows.jl
@@ -12,6 +12,9 @@ using Bijectors: PartitionMask, Inverse, combine, partition
 using Functors
 using AbstractPPL: AbstractPPL
 using LogExpFunctions: LogExpFunctions
+using ChainRulesCore: ignore_derivatives
+using GPUArraysCore: AbstractGPUMatrix
+using PDMats: PDMat, whiten
 
 using DocStringExtensions
 
@@ -135,6 +138,43 @@ function _device_specific_rand(
     return Random.rand(rng, td, n)
 end
 
+"""
+    _device_specific_logpdf(d, xs)
+
+Log-density of `d` at each column of `xs`, left on the device holding `xs`.
+`Distributions.logpdf` maps over the columns and materialises a host array, so the ELBO
+cannot be assembled from it when the samples live on a GPU.
+"""
+_device_specific_logpdf(d, xs::AbstractMatrix) = logpdf(d, xs)
+
+function _device_specific_logpdf(d::Distributions.MvNormal, xs::AbstractGPUMatrix)
+    return _batched_mvnormal_logpdf(d, xs)
+end
+
+# `logdet(::Cholesky)` accumulates `factors[i, i]` in a host loop, which a GPU array rejects.
+# Gathering the diagonal keeps it to one kernel. The other covariance types reduce over a
+# scalar or a vector in PDMats, so they need no help.
+_cov_logdet(Σ) = logdet(Σ)
+_cov_logdet(Σ::PDMat) = 2 * sum(log, diag(cholesky(Σ).factors))
+
+# Whole-array form of the multivariate normal log-density, so it runs wherever `xs` lives.
+# `whiten` stays on the device and does not mutate, unlike the `sqmahal` behind `logpdf`; a
+# solve against `d.Σ` would leave a `PDMats` tangent that AD cannot accumulate.
+#
+# `d` is held constant. Differentiating more than one use of a full covariance leaves a
+# cotangent per use, a `Diagonal` from the log-determinant and an `UpperTriangular` from the
+# whitening, and summing those two indexes a device array element by element. Base
+# distributions are leaves and targets are fixed, so no gradient is owed for `d` here, and
+# returning none beats returning a wrong one.
+function _batched_mvnormal_logpdf(d::Distributions.MvNormal, xs::AbstractMatrix)
+    T = eltype(xs)
+    μ = ignore_derivatives(d.μ)
+    Σ = ignore_derivatives(d.Σ)
+    c = ignore_derivatives(T(length(d) * log(2 * π)) + _cov_logdet(Σ))
+    q = sum(abs2, whiten(Σ, xs .- μ); dims=1)
+    return vec(-(c .+ q) ./ 2)
+end
+
 # interface of contructing common flow layers
 include("flows/utils.jl")
 include("flows/planar_radial.jl")
diff --git a/src/objectives/elbo.jl b/src/objectives/elbo.jl
index 6ddd47b8..9b87952a 100644
--- a/src/objectives/elbo.jl
+++ b/src/objectives/elbo.jl
@@ -65,7 +65,7 @@ Returns
 function _batched_elbos(flow::Bijectors.MultivariateTransformed, logp, xs::AbstractMatrix)
     # requires the flow transformation to be able to handle batched inputs
     ys, logabsdetjac = with_logabsdet_jacobian(flow.transform, xs) 
-    elbos = logp(ys) .- logpdf(flow.dist, xs) .+ logabsdetjac
+    elbos = logp(ys) .- _device_specific_logpdf(flow.dist, xs) .+ logabsdetjac
     return elbos
 end
 

From edde71b0322919a2bb12410e2c317bf0a0229815 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 03/17] Test the batched MvNormal log-density against
 Distributions

---
 test/objectives.jl | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/test/objectives.jl b/test/objectives.jl
index c8104726..8df79562 100644
--- a/test/objectives.jl
+++ b/test/objectives.jl
@@ -35,3 +35,34 @@
         end
     end
 end
+
+@testset "batched MvNormal log-density" begin
+    # `_batched_elbos` needs a log-density that stays on the sample's device, so the whole
+    # array form has to agree with Distributions on every covariance type.
+    @testset "$T" for T in (Float32, Float64)
+        # a non-unit scale, so dropping the whitening is caught in this case too
+        dists = (
+            MvNormal(zeros(T, 3), PDMats.ScalMat(3, T(2))),
+            MvNormal(T[1, -2, 0.5], Diagonal(T[2, 0.5, 1])),
+            MvNormal(zeros(T, 3), T[2 0.3 0.1; 0.3 1 0.2; 0.1 0.2 1.5]),
+        )
+        @testset "$(nameof(typeof(d.Σ)))" for d in dists
+            xs = randn(T, 3, 6)
+            # the gathered log-determinant has to agree with the one it replaces
+            @test NormalizingFlows._cov_logdet(d.Σ) ≈ logdet(d.Σ) rtol = sqrt(eps(T))
+
+            batched = NormalizingFlows._batched_mvnormal_logpdf(d, xs)
+            @test eltype(batched) == T
+            @test batched ≈ logpdf(d, xs) rtol = sqrt(eps(T))
+            # the generic fallback must leave the host path untouched
+            @test NormalizingFlows._device_specific_logpdf(d, xs) == logpdf(d, xs)
+
+            # it sits inside the differentiated ELBO, so the pullback has to work too
+            loss(v) = sum(NormalizingFlows._batched_mvnormal_logpdf(d, reshape(v, 3, 6)))
+            g_ref = ForwardDiff.gradient(loss, vec(xs))
+            @test all(isfinite, g_ref)
+            @test only(Zygote.gradient(loss, vec(xs))) ≈ g_ref rtol = sqrt(eps(T))
+            @test ReverseDiff.gradient(loss, vec(xs)) ≈ g_ref rtol = sqrt(eps(T))
+        end
+    end
+end

From c87a81ea678e4e28d7836f1974ef601949fc04f2 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 04/17] Test the log-density on device arrays without a GPU

---
 test/Project.toml |  6 +++++
 test/device.jl    | 57 +++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl  |  2 ++
 3 files changed, 65 insertions(+)
 create mode 100644 test/device.jl

diff --git a/test/Project.toml b/test/Project.toml
index af039cc8..93b83e59 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -7,10 +7,13 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -25,8 +28,11 @@ Enzyme = "0.13.186 - 0.13.188"
 Flux = "0.16.10"
 ForwardDiff = "1.4.1"
 Functors = "0.5.2"
+GPUArraysCore = "0.2"
+JLArrays = "0.3.3"
 Mooncake = "0.5"
 NormalizingFlows = "0.4.0"
 Optimisers = "0.4.7"
+PDMats = "0.11"
 ReverseDiff = "1.17.0"
 Zygote = "0.7.11"
diff --git a/test/device.jl b/test/device.jl
new file mode 100644
index 00000000..7b82de0b
--- /dev/null
+++ b/test/device.jl
@@ -0,0 +1,57 @@
+# JLArrays is the GPUArrays reference backend: it runs on the CPU but rejects scalar
+# indexing, so these tests reach device-only failures that would otherwise need a GPU.
+
+# Mirror what cuSOLVER returns: factors on the device with `uplo = 'U'`, which is what makes
+# PDMats wrap them in an `Adjoint` when whitening.
+function device_pdmat(A::AbstractMatrix)
+    c = cholesky(A)
+    return PDMats.PDMat(jl(A), Cholesky(jl(Matrix(c.factors)), c.uplo, c.info))
+end
+
+@testset "batched MvNormal log-density on device arrays" begin
+    GPUArraysCore.allowscalar(false)
+
+    @testset "$T" for T in (Float32, Float64)
+        rtol = T == Float32 ? 1.0f-4 : 1.0e-8
+        xs = randn(T, 2, 5)
+        covariances = (
+            PDMats.ScalMat(2, one(T)),
+            PDMats.PDiagMat(T[2, 0.5]),
+            PDMats.PDMat(T[2 0.3; 0.3 1]),
+        )
+
+        @testset "$(nameof(typeof(Σ)))" for Σ in covariances
+            host = MvNormal(zeros(T, 2), Σ)
+            Σ_dev = if Σ isa PDMats.PDMat
+                device_pdmat(Matrix(Σ))
+            elseif Σ isa PDMats.PDiagMat
+                PDMats.PDiagMat(jl(Σ.diag))
+            else
+                Σ
+            end
+            dev = MvNormal(jl(zeros(T, 2)), Σ_dev)
+            xs_dev = jl(xs)
+
+            batched = NormalizingFlows._batched_mvnormal_logpdf(dev, xs_dev)
+            @test batched isa JLArray
+            @test Array(batched) ≈ logpdf(host, xs) rtol = rtol
+
+            # the dispatch is on any GPU array, not just CUDA, so a second backend routes here
+            @test NormalizingFlows._device_specific_logpdf(dev, xs_dev) ≈ batched rtol =
+                rtol
+            # and the host path is still Distributions
+            @test NormalizingFlows._device_specific_logpdf(host, xs) == logpdf(host, xs)
+
+            # The gradient is the part that used to fail: a full covariance leaves one
+            # cotangent per differentiated use and adding them indexes the device array.
+            g = only(
+                Zygote.gradient(
+                    x -> sum(NormalizingFlows._batched_mvnormal_logpdf(dev, x)), xs_dev
+                ),
+            )
+            @test g isa JLArray
+            g_ref = ForwardDiff.gradient(x -> sum(logpdf(host, reshape(x, 2, 5))), vec(xs))
+            @test Array(g) ≈ reshape(g_ref, 2, 5) rtol = rtol
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 46e6137f..b2ce1938 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -6,6 +6,7 @@ using Random
 using ADTypes
 using Functors
 using ForwardDiff, Zygote, ReverseDiff, Enzyme, Mooncake
+using GPUArraysCore, JLArrays, PDMats
 using Flux: f32
 import DifferentiationInterface as DI
 
@@ -15,6 +16,7 @@ using Test
 @leaf MvNormal
 
 include("objectives.jl")
+include("device.jl")
 include("interface.jl")
 include("rqs.jl")
 include("flow.jl")

From bd1f9ec7c5ac04d78b3d945d824e094fe93d8719 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 05/17] Extend the CUDA tests to the batched ELBO and planar
 flow training

---
 test/ext/CUDA/Project.toml |   8 ++
 test/ext/CUDA/cuda.jl      | 146 +++++++++++++++++++++++++++++++------
 2 files changed, 132 insertions(+), 22 deletions(-)

diff --git a/test/ext/CUDA/Project.toml b/test/ext/CUDA/Project.toml
index 2a3e2f20..be9d3d57 100644
--- a/test/ext/CUDA/Project.toml
+++ b/test/ext/CUDA/Project.toml
@@ -1,18 +1,26 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+ADTypes = "1.22.2"
 Bijectors = "0.16.2"
 CUDA = "6.2.1"
+DifferentiationInterface = "0.7.20"
 Distributions = "0.25.129"
 Flux = "0.16.10"
+Functors = "0.5.2"
 NormalizingFlows = "0.4.0"
+Optimisers = "0.4.7"
 Zygote = "0.7.11"
diff --git a/test/ext/CUDA/cuda.jl b/test/ext/CUDA/cuda.jl
index 46944436..a3a5442c 100644
--- a/test/ext/CUDA/cuda.jl
+++ b/test/ext/CUDA/cuda.jl
@@ -3,33 +3,49 @@ Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
 using NormalizingFlows
-using Bijectors, CUDA, Distributions, Flux, LinearAlgebra, Random, Test
+using ADTypes, Bijectors, CUDA, Distributions, Flux, Functors, LinearAlgebra, Optimisers
+using Random, Test
 using Zygote
+# loads the AbstractPPL extension that routes `AutoZygote` through DifferentiationInterface
+import DifferentiationInterface as DI
+
+# keep q0 parameters out of Optimisers.destructure
+@leaf MvNormal
+
+# Bijectors' planar layer broadcasts in a way that CUDA cannot fuse, and reads `flow.b` back
+# from the device.
+# https://github.com/TuringLang/Bijectors.jl/blob/93cb25563043c527905519d81d6dee7917af4dbe/src/bijectors/planar_layer.jl#L65-L110
+function Bijectors.get_u_hat(u::CuVector{T}, w::CuVector{T}) where {T<:Real}
+    wT_u = dot(w, u)
+    scale = (Bijectors.LogExpFunctions.log1pexp(-wT_u) - 1) / sum(abs2, w)
+    û = CUDA.broadcast(+, u, CUDA.broadcast(*, scale, w))
+    wT_û = Bijectors.LogExpFunctions.log1pexp(wT_u) - 1
+    return û, wT_û
+end
+function Bijectors._transform(flow::PlanarLayer, z::CuArray{T}) where {T<:Real}
+    w = CuArray(flow.w)
 
-@testset "rand with CUDA" begin
-
-    # Bijectors versions use dot for broadcasting, which causes issues with CUDA.
-    # https://github.com/TuringLang/Bijectors.jl/blob/6f0d383f73afd150a018b65a3ea4ac9306065d38/src/bijectors/planar_layer.jl#L65-L80
-    function Bijectors.get_u_hat(u::CuVector{T}, w::CuVector{T}) where {T<:Real}
-        wT_u = dot(w, u)
-        scale = (Bijectors.LogExpFunctions.log1pexp(-wT_u) - 1) / sum(abs2, w)
-        û = CUDA.broadcast(+, u, CUDA.broadcast(*, scale, w))
-        wT_û = Bijectors.LogExpFunctions.log1pexp(wT_u) - 1
-        return û, wT_û
-    end
-    function Bijectors._transform(flow::PlanarLayer, z::CuArray{T}) where {T<:Real}
-        w = CuArray(flow.w)
-        b = T(first(flow.b))  # Scalar
-
-        û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
-        wT_z = Bijectors.aT_b(w, z)
+    û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
+    wT_z = Bijectors.aT_b(w, z)
 
-        tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, b))
-        transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
+    # `flow.b` holds one element, so broadcasting it is the same as Bijectors' `first(flow.b)`
+    # without reading back from the device.
+    tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, flow.b))
+    transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
 
-        return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
-    end
+    return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
+end
+# Only the batched path is covered. A `CuVector` still falls to Bijectors' own method, whose
+# `first(flow.b)` reads back from the device, so single samples need `allowscalar(true)`.
+function Bijectors.with_logabsdet_jacobian(
+    flow::PlanarLayer, z::CuMatrix{T}
+) where {T<:Real}
+    transformed, wT_û, wT_z = Bijectors._transform(flow, z)
+    logjac = log1p.(wT_û .* abs2.(sech.(vec(wT_z) .+ flow.b)))
+    return (result=transformed, logabsdetjac=logjac)
+end
 
+@testset "rand with CUDA" begin
     CUDA.allowscalar(true)
     dists = [
         MvNormal(CUDA.zeros(2), cu(Matrix{Float64}(I, 2, 2))),
@@ -98,3 +114,89 @@ end
     x_back, _ = NormalizingFlows.rqs_inverse(y_gpu, params_gpu...)
     @test Array(x_back) ≈ x_cpu rtol = 1.0f-4
 end
+
+# `Distributions.logpdf` maps over columns and returns a host array, which used to break the
+# batched ELBO on a GPU. Planar layers are used throughout because coupling layers still
+# partition through a host sparse `PartitionMask`.
+@testset "batched ELBO on CUDA" begin
+    CUDA.allowscalar(false)
+    q0 = MvNormal(CUDA.zeros(Float32, 2), cu(Matrix{Float32}(I, 2, 2)))
+    xs = NormalizingFlows._device_specific_rand(CUDA.default_rng(), q0, 64)
+
+    @testset "log-density stays on the device" begin
+        lp = NormalizingFlows._device_specific_logpdf(q0, xs)
+        @test lp isa CuArray{Float32}
+        @test length(lp) == 64
+        cpu_q0 = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
+        @test Array(lp) ≈ logpdf(cpu_q0, Array(xs)) rtol = 1.0f-4
+    end
+
+    @testset "batched ELBO" begin
+        target = MvNormal(CUDA.zeros(Float32, 2), cu(Matrix{Float32}(I, 2, 2)))
+        logp(z) = NormalizingFlows._device_specific_logpdf(target, z)
+        pl = PlanarLayer(
+            CUDA.rand(Float32, 2), CUDA.rand(Float32, 2), CUDA.rand(Float32, 1)
+        )
+        flow = Bijectors.transformed(q0, pl)
+
+        elbos = NormalizingFlows._batched_elbos(flow, logp, xs)
+        @test elbos isa CuArray{Float32}
+        @test all(isfinite, Array(elbos))
+        @test isfinite(elbo_batch(flow, logp, xs))
+
+        # The same flow on the host has to produce the same numbers, otherwise a wrong value
+        # in the planar overrides or the ELBO assembly would pass on types alone.
+        cpu_q0 = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
+        cpu_pl = PlanarLayer(Array(pl.u), Array(pl.w), Array(pl.b))
+        cpu_flow = Bijectors.transformed(cpu_q0, cpu_pl)
+        cpu_target = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
+        cpu_logp(z) = logpdf(cpu_target, z)
+        cpu_elbos = NormalizingFlows._batched_elbos(cpu_flow, cpu_logp, Array(xs))
+        @test Array(elbos) ≈ cpu_elbos rtol = 1.0f-4
+
+        # A full covariance has to differentiate, not just evaluate: each differentiated use
+        # of one leaves a cotangent of a different matrix type, and summing them indexes the
+        # device array element by element.
+        g = only(
+            Zygote.gradient(x -> sum(NormalizingFlows._batched_elbos(flow, logp, x)), xs)
+        )
+        @test g isa CuArray{Float32}
+        @test all(isfinite, Array(g))
+    end
+end
+
+@testset "planar flow training on CUDA" begin
+    CUDA.allowscalar(false)
+    T = Float32
+    d = 2
+
+    # Diagonal covariances exercise the broadcast branch of the log-density; the testset
+    # above covers the full covariance branch, gradient included.
+    q0 = MvNormal(CUDA.zeros(T, d), Diagonal(CUDA.ones(T, d)))
+    target = MvNormal(cu(T[2, -1]), Diagonal(CUDA.ones(T, d)))
+    logp(z) = NormalizingFlows._device_specific_logpdf(target, z)
+
+    layers = [PlanarLayer(CUDA.rand(T, d), CUDA.rand(T, d), CUDA.rand(T, 1)) for _ in 1:2]
+    flow = create_flow(layers, q0)
+
+    θ, re = Optimisers.destructure(flow)
+    @test θ isa CuArray{T}
+
+    flow_trained, stats, _ = train_flow(
+        CUDA.default_rng(),
+        elbo_batch,
+        flow,
+        logp,
+        32;
+        max_iters=5,
+        optimiser=Optimisers.Adam(T(1e-3)),
+        # Mooncake's CUDA support threads duals only through broadcast elements, so a
+        # broadcast whose closure captures differentiable values gets no gradient for them.
+        # Every layer here is such a broadcast.
+        ADbackend=ADTypes.AutoZygote(),
+        show_progress=false,
+    )
+
+    @test all(isfinite, map(x -> x.loss, stats))
+    @test Optimisers.destructure(flow_trained)[1] isa CuArray{T}
+end

From 683077eefd18040f27dc3e2ef527992ac0da3034 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 06/17] Add a GPU planar flow demo with its own project

---
 example/Project.toml     |   1 -
 example/README.md        |   2 +
 example/gpu/Project.toml |  22 +++++++
 example/gpu/demo_gpu.jl  |  87 +++++++++++++++++++++++++++
 example/test_gpu.jl      | 126 ---------------------------------------
 5 files changed, 111 insertions(+), 127 deletions(-)
 create mode 100644 example/gpu/Project.toml
 create mode 100644 example/gpu/demo_gpu.jl
 delete mode 100644 example/test_gpu.jl

diff --git a/example/Project.toml b/example/Project.toml
index 1645bed3..78de2230 100644
--- a/example/Project.toml
+++ b/example/Project.toml
@@ -2,7 +2,6 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
diff --git a/example/README.md b/example/README.md
index 738fff98..22515597 100644
--- a/example/README.md
+++ b/example/README.md
@@ -16,3 +16,5 @@ Currently, all examples share the same [Julia project](https://pkgdocs.julialang
 using Pkg; Pkg.activate("."); Pkg.instantiate()
 ```
 This will install all needed packages, at the exact versions when the model was last updated. Then you can run the model code with `include("<example-to-run>.jl")`, or by running the example script line-by-line.
+
+`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA. It needs a GPU; activate and instantiate `gpu/` to run it.
diff --git a/example/gpu/Project.toml b/example/gpu/Project.toml
new file mode 100644
index 00000000..321a590d
--- /dev/null
+++ b/example/gpu/Project.toml
@@ -0,0 +1,22 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+ADTypes = "1.22.2"
+Bijectors = "0.16.2"
+CUDA = "6.2.1"
+DifferentiationInterface = "0.7.20"
+Distributions = "0.25.129"
+Functors = "0.5.2"
+NormalizingFlows = "0.4.1"
+Optimisers = "0.4.7"
+Zygote = "0.7.11"
diff --git a/example/gpu/demo_gpu.jl b/example/gpu/demo_gpu.jl
new file mode 100644
index 00000000..70bde5b3
--- /dev/null
+++ b/example/gpu/demo_gpu.jl
@@ -0,0 +1,87 @@
+# Training a planar flow on the GPU. This demo has its own project so that the CPU examples
+# do not pull in CUDA. Run it from `example/gpu` after
+#
+#     using Pkg; Pkg.activate("."); Pkg.develop(; path="../.."); Pkg.instantiate()
+#
+# The `develop` is needed until 0.4.1 is registered, because the device log-density this
+# demo relies on landed in that version.
+#
+# Coupling flows (RealNVP, NSF) do not run on the GPU yet: `Bijectors.PartitionMask` holds
+# host sparse matrices and `partition`/`combine` multiply against them.
+
+using Distributions, LinearAlgebra
+using Bijectors
+using Functors
+using Optimisers, ADTypes, Zygote
+using CUDA
+using NormalizingFlows
+# loads the AbstractPPL extension that routes `AutoZygote` through DifferentiationInterface
+using DifferentiationInterface
+
+# Bijectors' planar layer broadcasts in a way that CUDA cannot fuse, and reads `flow.b` back
+# from the device.
+# https://github.com/TuringLang/Bijectors.jl/blob/93cb25563043c527905519d81d6dee7917af4dbe/src/bijectors/planar_layer.jl#L65-L110
+function Bijectors.get_u_hat(u::CuVector{T}, w::CuVector{T}) where {T<:Real}
+    wT_u = dot(w, u)
+    scale = (Bijectors.LogExpFunctions.log1pexp(-wT_u) - 1) / sum(abs2, w)
+    û = CUDA.broadcast(+, u, CUDA.broadcast(*, scale, w))
+    wT_û = Bijectors.LogExpFunctions.log1pexp(wT_u) - 1
+    return û, wT_û
+end
+function Bijectors._transform(flow::Bijectors.PlanarLayer, z::CuArray{T}) where {T<:Real}
+    w = CuArray(flow.w)
+    û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
+    wT_z = Bijectors.aT_b(w, z)
+    # `flow.b` holds one element, so broadcasting it is the same as Bijectors' `first(flow.b)`
+    # without reading back from the device.
+    tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, flow.b))
+    transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
+    return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
+end
+function Bijectors.with_logabsdet_jacobian(
+    flow::Bijectors.PlanarLayer, z::CuMatrix{T}
+) where {T<:Real}
+    transformed, wT_û, wT_z = Bijectors._transform(flow, z)
+    logjac = log1p.(wT_û .* abs2.(sech.(vec(wT_z) .+ flow.b)))
+    return (result=transformed, logabsdetjac=logjac)
+end
+
+rng = CUDA.default_rng()
+T = Float32
+d = 2
+
+@leaf MvNormal
+q0 = MvNormal(CUDA.zeros(T, d), Diagonal(CUDA.ones(T, d)))
+
+# `logp` takes the whole `(d, n)` batch and returns one value per column. Writing it with
+# array operations keeps it on the device, where `logpdf` would gather the columns onto the
+# host. The normaliser is kept so the reported ELBO is the true one.
+const μ_target = cu(T[2, -1])
+const logZ = T(d * log(2 * π))
+logp(z) = vec(-(logZ .+ sum(abs2, z .- μ_target; dims=1)) ./ 2)
+
+layers = [
+    Bijectors.PlanarLayer(CUDA.rand(T, d), CUDA.rand(T, d), CUDA.rand(T, 1)) for _ in 1:4
+]
+flow = create_flow(layers, q0)
+
+sample_per_iter = 64
+flow_trained, stats, _ = train_flow(
+    rng,
+    elbo_batch,
+    flow,
+    logp,
+    sample_per_iter;
+    max_iters=2_000,
+    optimiser=Optimisers.Adam(one(T) / 100),
+    # Mooncake's CUDA support threads duals only through broadcast elements, so a broadcast
+    # whose closure captures differentiable values gets no gradient for them. Every layer
+    # here is such a broadcast.
+    ADbackend=ADTypes.AutoZygote(),
+)
+
+losses = map(x -> x.loss, stats)
+@info "ELBO" start = -losses[1] final = -losses[end]
+
+ys = NormalizingFlows._device_specific_rand(rng, flow_trained, 1_000)
+@info "posterior mean (target is $(Array(μ_target)))" mean(Array(ys); dims=2)
diff --git a/example/test_gpu.jl b/example/test_gpu.jl
deleted file mode 100644
index f6da860d..00000000
--- a/example/test_gpu.jl
+++ /dev/null
@@ -1,126 +0,0 @@
-using Flux
-using Bijectors
-using Bijectors: partition, combine, PartitionMask
-
-using Random, Distributions, LinearAlgebra
-using Functors
-using Optimisers, ADTypes
-using Mooncake
-using NormalizingFlows
-
-include("SyntheticTargets.jl")
-include("utils.jl")
-
-##################################
-# define affine coupling layer using Bijectors.jl interface
-#################################
-struct AffineCoupling <: Bijectors.Bijector
-    dim::Int
-    mask::Bijectors.PartitionMask
-    s::Flux.Chain
-    t::Flux.Chain
-end
-
-# let params track field s and t
-@functor AffineCoupling (s, t)
-
-function AffineCoupling(
-    dim::Int,  # dimension of input
-    hdims::Int, # dimension of hidden units for s and t
-    mask_idx::AbstractVector, # index of dimensione that one wants to apply transformations on
-)
-    cdims = length(mask_idx) # dimension of parts used to construct coupling law
-    s = mlp3(cdims, hdims, cdims)
-    t = mlp3(cdims, hdims, cdims)
-    mask = PartitionMask(dim, mask_idx)
-    return AffineCoupling(dim, mask, s, t)
-end
-
-function Bijectors.transform(af::AffineCoupling, x::AbstractVecOrMat)
-    # partition vector using 'af.mask::PartitionMask`
-    x₁, x₂, x₃ = partition(af.mask, x)
-    y₁ = x₁ .* af.s(x₂) .+ af.t(x₂)
-    return combine(af.mask, y₁, x₂, x₃)
-end
-
-function (af::AffineCoupling)(x::AbstractArray)
-    return transform(af, x)
-end
-
-function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractVector)
-    x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
-    y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
-    logjac = sum(log ∘ abs, af.s(x_2)) # this is a scalar
-    return combine(af.mask, y_1, x_2, x_3), logjac
-end
-
-function Bijectors.with_logabsdet_jacobian(af::AffineCoupling, x::AbstractMatrix)
-    x_1, x_2, x_3 = Bijectors.partition(af.mask, x)
-    y_1 = af.s(x_2) .* x_1 .+ af.t(x_2)
-    logjac = sum(log ∘ abs, af.s(x_2); dims = 1) # 1 × size(x, 2)
-    return combine(af.mask, y_1, x_2, x_3), vec(logjac)
-end
-
-
-function Bijectors.with_logabsdet_jacobian(
-    iaf::Inverse{<:AffineCoupling}, y::AbstractVector
-)
-    af = iaf.orig
-    # partition vector using `af.mask::PartitionMask`
-    y_1, y_2, y_3 = partition(af.mask, y)
-    # inverse transformation
-    x_1 = (y_1 .- af.t(y_2)) ./ af.s(y_2)
-    logjac = -sum(log ∘ abs, af.s(y_2))
-    return combine(af.mask, x_1, y_2, y_3), logjac
-end
-
-function Bijectors.with_logabsdet_jacobian(
-    iaf::Inverse{<:AffineCoupling}, y::AbstractMatrix
-)
-    af = iaf.orig
-    # partition vector using `af.mask::PartitionMask`
-    y_1, y_2, y_3 = partition(af.mask, y)
-    # inverse transformation
-    x_1 = (y_1 .- af.t(y_2)) ./ af.s(y_2)
-    logjac = -sum(log ∘ abs, af.s(y_2); dims = 1)
-    return combine(af.mask, x_1, y_2, y_3), vec(logjac)
-end
-
-
-##################################
-# start demo
-#################################
-using CUDA
-const NF = NormalizingFlows
-rng_g = CUDA.default_rng() # use GPU RNG if available
-
-
-CUDA.allowscalar(true)
-n_samples = 100
-q0 = MvNormal(CUDA.zeros(2), cu([1f0 0f0; 0f0 1f0]))
-# gpu sample from the reference
-xs = NF._device_specific_rand(rng_g, q0, n_samples)
-
-d = 2
-hdims = 32
-Ls_g = [AffineCoupling(d, hdims, [1]) ∘ AffineCoupling(d, hdims, [2]) for i in 1:3]
-flow_g = create_flow(Ls_g, q0)
-flow_g = fmap(cu, flow_g) # move all flow parameters be on GPU
-
-# gpu sample from the flow
-ys = NF._device_specific_rand(rng_g, flow_g, n_samples)
-
-# log density computation
-logpdf(flow_g, ys) # errored
-
-logpdf(q0, xs) # returns a CPU array
-
-
-# elbo_batch(rng_g, flow, logp, n_samples)
-
-target = Banana(2, 1.0f0, 100.0f0)
-target_g = fmap(cu, target) # move target to GPU
-logp_g = Base.Fix1(logpdf, target_g)
-
-logp_g(yy)
-

From a8a9cac1c92ed0a1bdbb628d867a10fcf2c31973 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 07/17] Correct the nsf note about running on the GPU

---
 src/flows/neuralspline.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/flows/neuralspline.jl b/src/flows/neuralspline.jl
index 8b799921..6fad16c1 100644
--- a/src/flows/neuralspline.jl
+++ b/src/flows/neuralspline.jl
@@ -226,9 +226,11 @@ Returns
 - `Bijectors.TransformedDistribution` representing the NSF flow.
 
 !!! note
-    The rational quadratic spline is written with whole-array operations, so the flow runs
-    on the CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`,
-    `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`.
+    The rational quadratic spline is written with whole-array operations, so it runs on the
+    CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`,
+    `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`. The assembled flow is CPU only: the
+    coupling layer splits its input through `Bijectors.PartitionMask`, which holds host
+    sparse matrices.
 
 !!! note
     When training the flow, mark the base distribution as a leaf first (`Functors.@leaf

From 0f53834cf1852a5970a8897ca255eea7b98884e7 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 08/17] Add a changelog entry for the GPU batched ELBO

---
 HISTORY.md | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/HISTORY.md b/HISTORY.md
index 5c9474df..bfe58b1a 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,3 +1,19 @@
+# 0.4.1
+
+## Other changes
+
+`elbo_batch` now works when the samples live on a GPU.
+It used to assemble the ELBO from `Distributions.logpdf`, which evaluates the base distribution column by column and returns a host array.
+The base distribution's log-density is now taken with a whole-array form that stays on the device holding the samples, for any GPU array backend.
+`elbo` is unchanged and still goes through `Distributions.logpdf`, so it stays CPU only.
+
+Coupling flows (`realnvp`, `nsf`) still do not run on the GPU.
+`Bijectors.PartitionMask` holds host sparse matrices and `partition`/`combine` multiply against them, so the split is done on the host.
+`example/gpu/demo_gpu.jl` trains a planar flow instead.
+
+The base distribution is treated as a constant when its log-density is taken this way.
+Differentiating more than one use of a full covariance leaves a cotangent per use, and summing those indexes a device array element by element.
+
 # 0.4.0
 
 ## Breaking changes

From 2e35568102e23ea4462ee7dd3d41e7b98c8edbd8 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:25:48 +0530
Subject: [PATCH 09/17] Bump patch version to 0.4.1

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index a8836a82..10b2194c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "NormalizingFlows"
 uuid = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

From 1943d7753be7b1d390a62583b6a7ac1931e816f7 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:39:38 +0530
Subject: [PATCH 10/17] Move the planar layer to the host structurally in the
 CUDA reference check

---
 test/ext/CUDA/cuda.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/ext/CUDA/cuda.jl b/test/ext/CUDA/cuda.jl
index a3a5442c..dc70b1a1 100644
--- a/test/ext/CUDA/cuda.jl
+++ b/test/ext/CUDA/cuda.jl
@@ -147,7 +147,7 @@ end
         # The same flow on the host has to produce the same numbers, otherwise a wrong value
         # in the planar overrides or the ELBO assembly would pass on types alone.
         cpu_q0 = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
-        cpu_pl = PlanarLayer(Array(pl.u), Array(pl.w), Array(pl.b))
+        cpu_pl = fmap(Array, pl)
         cpu_flow = Bijectors.transformed(cpu_q0, cpu_pl)
         cpu_target = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
         cpu_logp(z) = logpdf(cpu_target, z)

From d190e520d2a8ecee3dc6d775434b4945f321e5c6 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Fri, 28 Aug 2026 14:55:36 +0530
Subject: [PATCH 11/17] State the observed Mooncake CUDA failure instead of the
 assumed cause

---
 example/gpu/demo_gpu.jl | 5 ++---
 test/ext/CUDA/cuda.jl   | 5 ++---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/example/gpu/demo_gpu.jl b/example/gpu/demo_gpu.jl
index 70bde5b3..510473ad 100644
--- a/example/gpu/demo_gpu.jl
+++ b/example/gpu/demo_gpu.jl
@@ -74,9 +74,8 @@ flow_trained, stats, _ = train_flow(
     sample_per_iter;
     max_iters=2_000,
     optimiser=Optimisers.Adam(one(T) / 100),
-    # Mooncake's CUDA support threads duals only through broadcast elements, so a broadcast
-    # whose closure captures differentiable values gets no gradient for them. Every layer
-    # here is such a broadcast.
+    # Zygote rather than Mooncake: Mooncake 0.5.48 fails on this flow inside its own CUDA
+    # kernel launch, a `CoDual` type assertion in `Adapt.adapt_storage`.
     ADbackend=ADTypes.AutoZygote(),
 )
 
diff --git a/test/ext/CUDA/cuda.jl b/test/ext/CUDA/cuda.jl
index dc70b1a1..d534eaa6 100644
--- a/test/ext/CUDA/cuda.jl
+++ b/test/ext/CUDA/cuda.jl
@@ -190,9 +190,8 @@ end
         32;
         max_iters=5,
         optimiser=Optimisers.Adam(T(1e-3)),
-        # Mooncake's CUDA support threads duals only through broadcast elements, so a
-        # broadcast whose closure captures differentiable values gets no gradient for them.
-        # Every layer here is such a broadcast.
+        # Zygote rather than Mooncake: Mooncake 0.5.48 fails on this flow inside its own
+        # CUDA kernel launch, a `CoDual` type assertion in `Adapt.adapt_storage`.
         ADbackend=ADTypes.AutoZygote(),
         show_progress=false,
     )

From 1805b2af66bf21c90d72c9a600bf7de17f98e9c5 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 01:24:29 +0530
Subject: [PATCH 12/17] Give Mooncake its own zero-derivative rule for the
 device draw

---
 Project.toml                       |  3 +++
 ext/NormalizingFlowsCUDAExt.jl     | 16 ++++++----------
 ext/NormalizingFlowsMooncakeExt.jl | 12 ++++++++++++
 src/NormalizingFlows.jl            | 15 ++++++++++++++-
 4 files changed, 35 insertions(+), 11 deletions(-)
 create mode 100644 ext/NormalizingFlowsMooncakeExt.jl

diff --git a/Project.toml b/Project.toml
index 10b2194c..2c477138 100644
--- a/Project.toml
+++ b/Project.toml
@@ -22,10 +22,12 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
 NormalizingFlowsCUDAExt = "CUDA"
+NormalizingFlowsMooncakeExt = "Mooncake"
 NormalizingFlowsReverseDiffExt = "ReverseDiff"
 
 [compat]
@@ -40,6 +42,7 @@ Flux = "0.16"
 Functors = "0.5.2"
 GPUArraysCore = "0.1, 0.2"
 LogExpFunctions = "0.3.3"
+Mooncake = "0.5"
 Optimisers = "0.2.16, 0.3, 0.4"
 PDMats = "0.11"
 ProgressMeter = "1.0.0"
diff --git a/ext/NormalizingFlowsCUDAExt.jl b/ext/NormalizingFlowsCUDAExt.jl
index b7d7e0c0..ed75c0bd 100644
--- a/ext/NormalizingFlowsCUDAExt.jl
+++ b/ext/NormalizingFlowsCUDAExt.jl
@@ -3,7 +3,6 @@ module NormalizingFlowsCUDAExt
 using CUDA
 using NormalizingFlows
 using NormalizingFlows: Bijectors, Distributions, Random
-using ChainRulesCore: @non_differentiable
 
 function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG,
@@ -24,7 +23,7 @@ function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
-    return _cuda_draw(rng, s, size(s))
+    return NormalizingFlows._device_draw(rng, s, size(s))
 end
 
 function _cuda_rand(
@@ -32,20 +31,15 @@ function _cuda_rand(
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
 )
-    return _cuda_draw(rng, s, (size(s)..., n))
+    return NormalizingFlows._device_draw(rng, s, (size(s)..., n))
 end
 
-function _cuda_draw(rng::CUDA.RNG, s, dims::Tuple)
+function NormalizingFlows._device_draw(rng::CUDA.RNG, s, dims::Tuple)
     return @inbounds Distributions.rand!(
         rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, dims)
     )
 end
 
-# Zygote cannot trace a `CuArray` allocation: it descends into the CUDA allocator and fails
-# to compile. Nothing is lost by hiding the draw, since `rand!` carries no gradient on the
-# host path either. Only the draw is opaque; a flow's transform stays differentiable.
-@non_differentiable _cuda_draw(::Any, ::Any, ::Any)
-
 # ! this is type piracy
 # replacing original function with scalar indexing
 function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVecOrMat)
@@ -56,7 +50,9 @@ function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVec
 end
 
 # to enable `_device_specific_rand(rng:CUDA.RNG, flow[, num_samples])`
-function NormalizingFlows._device_specific_rand(rng::CUDA.RNG, td::Bijectors.TransformedDistribution)
+function NormalizingFlows._device_specific_rand(
+    rng::CUDA.RNG, td::Bijectors.TransformedDistribution
+)
     return _cuda_rand(rng, td)
 end
 
diff --git a/ext/NormalizingFlowsMooncakeExt.jl b/ext/NormalizingFlowsMooncakeExt.jl
new file mode 100644
index 00000000..75aff1b0
--- /dev/null
+++ b/ext/NormalizingFlowsMooncakeExt.jl
@@ -0,0 +1,12 @@
+module NormalizingFlowsMooncakeExt
+
+using Mooncake: Mooncake, DefaultCtx
+using NormalizingFlows: NormalizingFlows
+
+# Mooncake imports ChainRules rules one signature at a time, so it does not see the
+# `@non_differentiable` on the device draw and traces into the allocation instead.
+Mooncake.@zero_derivative DefaultCtx Tuple{
+    typeof(NormalizingFlows._device_draw),Any,Any,Any
+}
+
+end
diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl
index 59043685..046a0a40 100644
--- a/src/NormalizingFlows.jl
+++ b/src/NormalizingFlows.jl
@@ -12,7 +12,7 @@ using Bijectors: PartitionMask, Inverse, combine, partition
 using Functors
 using AbstractPPL: AbstractPPL
 using LogExpFunctions: LogExpFunctions
-using ChainRulesCore: ignore_derivatives
+using ChainRulesCore: @non_differentiable, ignore_derivatives
 using GPUArraysCore: AbstractGPUMatrix
 using PDMats: PDMat, whiten
 
@@ -138,6 +138,19 @@ function _device_specific_rand(
     return Random.rand(rng, td, n)
 end
 
+"""
+    _device_draw(rng, s, dims)
+
+Draw a `dims`-shaped sample from `s` into an array on the device `rng` targets. Device
+extensions add the methods.
+"""
+function _device_draw end
+
+# The draw carries no gradient, and no AD backend can trace a device allocation: Zygote
+# descends into the allocator and fails to compile. Mooncake reads ChainRules rules only one
+# signature at a time, so it needs its own declaration, which the Mooncake extension adds.
+@non_differentiable _device_draw(::Any, ::Any, ::Any)
+
 """
     _device_specific_logpdf(d, xs)
 

From 4f094f4ecc15c9970ae047efb14065aae66d3681 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 01:24:29 +0530
Subject: [PATCH 13/17] Train with Mooncake as well as Zygote in the CUDA tests

---
 test/ext/CUDA/Project.toml |  2 ++
 test/ext/CUDA/cuda.jl      | 55 ++++++++++++++++++++++----------------
 2 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/test/ext/CUDA/Project.toml b/test/ext/CUDA/Project.toml
index be9d3d57..ad06b88f 100644
--- a/test/ext/CUDA/Project.toml
+++ b/test/ext/CUDA/Project.toml
@@ -7,6 +7,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -21,6 +22,7 @@ DifferentiationInterface = "0.7.20"
 Distributions = "0.25.129"
 Flux = "0.16.10"
 Functors = "0.5.2"
+Mooncake = "0.5.48"
 NormalizingFlows = "0.4.0"
 Optimisers = "0.4.7"
 Zygote = "0.7.11"
diff --git a/test/ext/CUDA/cuda.jl b/test/ext/CUDA/cuda.jl
index d534eaa6..d1245f1f 100644
--- a/test/ext/CUDA/cuda.jl
+++ b/test/ext/CUDA/cuda.jl
@@ -5,7 +5,7 @@ Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 using NormalizingFlows
 using ADTypes, Bijectors, CUDA, Distributions, Flux, Functors, LinearAlgebra, Optimisers
 using Random, Test
-using Zygote
+using Mooncake, Zygote
 # loads the AbstractPPL extension that routes `AutoZygote` through DifferentiationInterface
 import DifferentiationInterface as DI
 
@@ -176,26 +176,35 @@ end
     target = MvNormal(cu(T[2, -1]), Diagonal(CUDA.ones(T, d)))
     logp(z) = NormalizingFlows._device_specific_logpdf(target, z)
 
-    layers = [PlanarLayer(CUDA.rand(T, d), CUDA.rand(T, d), CUDA.rand(T, 1)) for _ in 1:2]
-    flow = create_flow(layers, q0)
-
-    θ, re = Optimisers.destructure(flow)
-    @test θ isa CuArray{T}
-
-    flow_trained, stats, _ = train_flow(
-        CUDA.default_rng(),
-        elbo_batch,
-        flow,
-        logp,
-        32;
-        max_iters=5,
-        optimiser=Optimisers.Adam(T(1e-3)),
-        # Zygote rather than Mooncake: Mooncake 0.5.48 fails on this flow inside its own
-        # CUDA kernel launch, a `CoDual` type assertion in `Adapt.adapt_storage`.
-        ADbackend=ADTypes.AutoZygote(),
-        show_progress=false,
-    )
-
-    @test all(isfinite, map(x -> x.loss, stats))
-    @test Optimisers.destructure(flow_trained)[1] isa CuArray{T}
+    backends = ADTypes.AbstractADType[ADTypes.AutoZygote()]
+    # On Julia 1.10 the broadcast kernel reaches Mooncake as a KernelAbstractions foreign
+    # call, which it does not differentiate.
+    if VERSION >= v"1.11"
+        push!(backends, ADTypes.AutoMooncake(; config=Mooncake.Config()))
+    end
+
+    @testset "$(nameof(typeof(ad)))" for ad in backends
+        layers = [
+            PlanarLayer(CUDA.rand(T, d), CUDA.rand(T, d), CUDA.rand(T, 1)) for _ in 1:2
+        ]
+        flow = create_flow(layers, q0)
+
+        θ, re = Optimisers.destructure(flow)
+        @test θ isa CuArray{T}
+
+        flow_trained, stats, _ = train_flow(
+            CUDA.default_rng(),
+            elbo_batch,
+            flow,
+            logp,
+            32;
+            max_iters=5,
+            optimiser=Optimisers.Adam(T(1e-3)),
+            ADbackend=ad,
+            show_progress=false,
+        )
+
+        @test all(isfinite, map(x -> x.loss, stats))
+        @test Optimisers.destructure(flow_trained)[1] isa CuArray{T}
+    end
 end

From 523f113a55026e7adae1495aa130fc33284e833d Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 01:24:29 +0530
Subject: [PATCH 14/17] Use Mooncake in the GPU demo

---
 HISTORY.md               |  1 +
 example/gpu/Project.toml |  6 ++----
 example/gpu/demo_gpu.jl  | 10 ++++------
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index bfe58b1a..5e8f06ca 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -6,6 +6,7 @@
 It used to assemble the ELBO from `Distributions.logpdf`, which evaluates the base distribution column by column and returns a host array.
 The base distribution's log-density is now taken with a whole-array form that stays on the device holding the samples, for any GPU array backend.
 `elbo` is unchanged and still goes through `Distributions.logpdf`, so it stays CPU only.
+Mooncake now differentiates the GPU path: it reads ChainRules rules one signature at a time, so the device draw carries a `Mooncake.@zero_derivative` declaration of its own in a new extension.
 
 Coupling flows (`realnvp`, `nsf`) still do not run on the GPU.
 `Bijectors.PartitionMask` holds host sparse matrices and `partition`/`combine` multiply against them, so the split is done on the host.
diff --git a/example/gpu/Project.toml b/example/gpu/Project.toml
index 321a590d..fc808a1c 100644
--- a/example/gpu/Project.toml
+++ b/example/gpu/Project.toml
@@ -2,21 +2,19 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ADTypes = "1.22.2"
 Bijectors = "0.16.2"
 CUDA = "6.2.1"
-DifferentiationInterface = "0.7.20"
 Distributions = "0.25.129"
 Functors = "0.5.2"
+Mooncake = "0.5.48"
 NormalizingFlows = "0.4.1"
 Optimisers = "0.4.7"
-Zygote = "0.7.11"
diff --git a/example/gpu/demo_gpu.jl b/example/gpu/demo_gpu.jl
index 510473ad..7d9d7070 100644
--- a/example/gpu/demo_gpu.jl
+++ b/example/gpu/demo_gpu.jl
@@ -12,11 +12,9 @@
 using Distributions, LinearAlgebra
 using Bijectors
 using Functors
-using Optimisers, ADTypes, Zygote
+using Optimisers, ADTypes, Mooncake
 using CUDA
 using NormalizingFlows
-# loads the AbstractPPL extension that routes `AutoZygote` through DifferentiationInterface
-using DifferentiationInterface
 
 # Bijectors' planar layer broadcasts in a way that CUDA cannot fuse, and reads `flow.b` back
 # from the device.
@@ -74,9 +72,9 @@ flow_trained, stats, _ = train_flow(
     sample_per_iter;
     max_iters=2_000,
     optimiser=Optimisers.Adam(one(T) / 100),
-    # Zygote rather than Mooncake: Mooncake 0.5.48 fails on this flow inside its own CUDA
-    # kernel launch, a `CoDual` type assertion in `Adapt.adapt_storage`.
-    ADbackend=ADTypes.AutoZygote(),
+    # Needs Julia 1.11 or newer: on 1.10 the broadcast kernel reaches Mooncake as a
+    # KernelAbstractions foreign call, which it does not differentiate.
+    ADbackend=ADTypes.AutoMooncake(; config=Mooncake.Config()),
 )
 
 losses = map(x -> x.loss, stats)

From 122b49f173e6774c0f9d3bc052df7efcdd224e3e Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 02:10:16 +0530
Subject: [PATCH 15/17] Sum the mapped array so the GPU rule set covers it

---
 src/NormalizingFlows.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl
index 046a0a40..4f880ee8 100644
--- a/src/NormalizingFlows.jl
+++ b/src/NormalizingFlows.jl
@@ -184,7 +184,9 @@ function _batched_mvnormal_logpdf(d::Distributions.MvNormal, xs::AbstractMatrix)
     μ = ignore_derivatives(d.μ)
     Σ = ignore_derivatives(d.Σ)
     c = ignore_derivatives(T(length(d) * log(2 * π)) + _cov_logdet(Σ))
-    q = sum(abs2, whiten(Σ, xs .- μ); dims=1)
+    # `sum(f, x; dims)` has no GPU rule in Mooncake, but the mapped array and `sum(x; dims)`
+    # both do.
+    q = sum(abs2.(whiten(Σ, xs .- μ)); dims=1)
     return vec(-(c .+ q) ./ 2)
 end
 

From 2228fe565f6d360db241cc2821a6a862a670c2db Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 21:30:37 +0530
Subject: [PATCH 16/17] Trim the comments and changelog to the non-obvious why

---
 HISTORY.md                         | 17 +++++------------
 example/README.md                  |  2 +-
 example/gpu/demo_gpu.jl            | 17 +++++------------
 ext/NormalizingFlowsMooncakeExt.jl |  3 +--
 src/NormalizingFlows.jl            | 24 +++++++-----------------
 test/device.jl                     | 13 +++++--------
 test/ext/CUDA/cuda.jl              | 21 +++++++--------------
 test/objectives.jl                 |  5 -----
 8 files changed, 31 insertions(+), 71 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index 5e8f06ca..b7395256 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -2,18 +2,11 @@
 
 ## Other changes
 
-`elbo_batch` now works when the samples live on a GPU.
-It used to assemble the ELBO from `Distributions.logpdf`, which evaluates the base distribution column by column and returns a host array.
-The base distribution's log-density is now taken with a whole-array form that stays on the device holding the samples, for any GPU array backend.
-`elbo` is unchanged and still goes through `Distributions.logpdf`, so it stays CPU only.
-Mooncake now differentiates the GPU path: it reads ChainRules rules one signature at a time, so the device draw carries a `Mooncake.@zero_derivative` declaration of its own in a new extension.
-
-Coupling flows (`realnvp`, `nsf`) still do not run on the GPU.
-`Bijectors.PartitionMask` holds host sparse matrices and `partition`/`combine` multiply against them, so the split is done on the host.
-`example/gpu/demo_gpu.jl` trains a planar flow instead.
-
-The base distribution is treated as a constant when its log-density is taken this way.
-Differentiating more than one use of a full covariance leaves a cotangent per use, and summing those indexes a device array element by element.
+`elbo_batch` works when the samples live on a GPU, for any GPU array backend.
+The base distribution's log-density is taken with a whole-array form that stays on the device, and the distribution is held constant under automatic differentiation.
+Zygote and Mooncake both differentiate the GPU path, Mooncake through a new extension.
+`elbo` and the coupling flows stay on the host, the latter because `Bijectors.PartitionMask` holds host sparse matrices.
+`example/gpu/demo_gpu.jl` trains a planar flow on the GPU.
 
 # 0.4.0
 
diff --git a/example/README.md b/example/README.md
index 22515597..66045714 100644
--- a/example/README.md
+++ b/example/README.md
@@ -17,4 +17,4 @@ using Pkg; Pkg.activate("."); Pkg.instantiate()
 ```
 This will install all needed packages, at the exact versions when the model was last updated. Then you can run the model code with `include("<example-to-run>.jl")`, or by running the example script line-by-line.
 
-`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA. It needs a GPU; activate and instantiate `gpu/` to run it.
+`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA. Activate and instantiate `gpu/` to run it.
diff --git a/example/gpu/demo_gpu.jl b/example/gpu/demo_gpu.jl
index 7d9d7070..a144257a 100644
--- a/example/gpu/demo_gpu.jl
+++ b/example/gpu/demo_gpu.jl
@@ -1,13 +1,8 @@
-# Training a planar flow on the GPU. This demo has its own project so that the CPU examples
-# do not pull in CUDA. Run it from `example/gpu` after
+# Training a planar flow on the GPU. This demo has its own project, run it from `example/gpu`:
 #
 #     using Pkg; Pkg.activate("."); Pkg.develop(; path="../.."); Pkg.instantiate()
 #
-# The `develop` is needed until 0.4.1 is registered, because the device log-density this
-# demo relies on landed in that version.
-#
-# Coupling flows (RealNVP, NSF) do not run on the GPU yet: `Bijectors.PartitionMask` holds
-# host sparse matrices and `partition`/`combine` multiply against them.
+# The `develop` is needed until 0.4.1 is registered.
 
 using Distributions, LinearAlgebra
 using Bijectors
@@ -30,8 +25,7 @@ function Bijectors._transform(flow::Bijectors.PlanarLayer, z::CuArray{T}) where
     w = CuArray(flow.w)
     û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
     wT_z = Bijectors.aT_b(w, z)
-    # `flow.b` holds one element, so broadcasting it is the same as Bijectors' `first(flow.b)`
-    # without reading back from the device.
+    # `flow.b` holds one element, so broadcasting it avoids reading back from the device.
     tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, flow.b))
     transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
     return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
@@ -51,9 +45,8 @@ d = 2
 @leaf MvNormal
 q0 = MvNormal(CUDA.zeros(T, d), Diagonal(CUDA.ones(T, d)))
 
-# `logp` takes the whole `(d, n)` batch and returns one value per column. Writing it with
-# array operations keeps it on the device, where `logpdf` would gather the columns onto the
-# host. The normaliser is kept so the reported ELBO is the true one.
+# `logp` takes the whole batch and returns one value per column. Array operations keep it on
+# the device, where `logpdf` would gather the columns onto the host.
 const μ_target = cu(T[2, -1])
 const logZ = T(d * log(2 * π))
 logp(z) = vec(-(logZ .+ sum(abs2, z .- μ_target; dims=1)) ./ 2)
diff --git a/ext/NormalizingFlowsMooncakeExt.jl b/ext/NormalizingFlowsMooncakeExt.jl
index 75aff1b0..b08889f0 100644
--- a/ext/NormalizingFlowsMooncakeExt.jl
+++ b/ext/NormalizingFlowsMooncakeExt.jl
@@ -3,8 +3,7 @@ module NormalizingFlowsMooncakeExt
 using Mooncake: Mooncake, DefaultCtx
 using NormalizingFlows: NormalizingFlows
 
-# Mooncake imports ChainRules rules one signature at a time, so it does not see the
-# `@non_differentiable` on the device draw and traces into the allocation instead.
+# Mooncake does not read the ChainRules rule on `_device_draw`.
 Mooncake.@zero_derivative DefaultCtx Tuple{
     typeof(NormalizingFlows._device_draw),Any,Any,Any
 }
diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl
index 4f880ee8..e327f827 100644
--- a/src/NormalizingFlows.jl
+++ b/src/NormalizingFlows.jl
@@ -146,9 +146,8 @@ extensions add the methods.
 """
 function _device_draw end
 
-# The draw carries no gradient, and no AD backend can trace a device allocation: Zygote
-# descends into the allocator and fails to compile. Mooncake reads ChainRules rules only one
-# signature at a time, so it needs its own declaration, which the Mooncake extension adds.
+# No AD backend can trace a device allocation. Mooncake reads ChainRules rules one signature
+# at a time, so it needs the separate declaration in its extension.
 @non_differentiable _device_draw(::Any, ::Any, ::Any)
 
 """
@@ -164,28 +163,19 @@ function _device_specific_logpdf(d::Distributions.MvNormal, xs::AbstractGPUMatri
     return _batched_mvnormal_logpdf(d, xs)
 end
 
-# `logdet(::Cholesky)` accumulates `factors[i, i]` in a host loop, which a GPU array rejects.
-# Gathering the diagonal keeps it to one kernel. The other covariance types reduce over a
-# scalar or a vector in PDMats, so they need no help.
+# `logdet(::Cholesky)` reads `factors[i, i]` in a host loop, which a GPU array rejects.
 _cov_logdet(Σ) = logdet(Σ)
 _cov_logdet(Σ::PDMat) = 2 * sum(log, diag(cholesky(Σ).factors))
 
-# Whole-array form of the multivariate normal log-density, so it runs wherever `xs` lives.
-# `whiten` stays on the device and does not mutate, unlike the `sqmahal` behind `logpdf`; a
-# solve against `d.Σ` would leave a `PDMats` tangent that AD cannot accumulate.
-#
-# `d` is held constant. Differentiating more than one use of a full covariance leaves a
-# cotangent per use, a `Diagonal` from the log-determinant and an `UpperTriangular` from the
-# whitening, and summing those two indexes a device array element by element. Base
-# distributions are leaves and targets are fixed, so no gradient is owed for `d` here, and
-# returning none beats returning a wrong one.
+# `whiten` stays on the device and does not mutate, unlike the `sqmahal` behind `logpdf`.
+# `d` is held constant because differentiating two uses of a full covariance leaves a
+# cotangent per use, and summing those indexes a device array element by element.
 function _batched_mvnormal_logpdf(d::Distributions.MvNormal, xs::AbstractMatrix)
     T = eltype(xs)
     μ = ignore_derivatives(d.μ)
     Σ = ignore_derivatives(d.Σ)
     c = ignore_derivatives(T(length(d) * log(2 * π)) + _cov_logdet(Σ))
-    # `sum(f, x; dims)` has no GPU rule in Mooncake, but the mapped array and `sum(x; dims)`
-    # both do.
+    # Mooncake has no GPU rule for `sum(f, x; dims)`.
     q = sum(abs2.(whiten(Σ, xs .- μ)); dims=1)
     return vec(-(c .+ q) ./ 2)
 end
diff --git a/test/device.jl b/test/device.jl
index 7b82de0b..8b78f824 100644
--- a/test/device.jl
+++ b/test/device.jl
@@ -1,8 +1,7 @@
-# JLArrays is the GPUArrays reference backend: it runs on the CPU but rejects scalar
-# indexing, so these tests reach device-only failures that would otherwise need a GPU.
+# JLArrays runs on the CPU but rejects scalar indexing, so it reaches device-only failures
+# without a GPU.
 
-# Mirror what cuSOLVER returns: factors on the device with `uplo = 'U'`, which is what makes
-# PDMats wrap them in an `Adjoint` when whitening.
+# cuSOLVER returns `uplo = 'U'`, which is what makes PDMats wrap the factors in an `Adjoint`.
 function device_pdmat(A::AbstractMatrix)
     c = cholesky(A)
     return PDMats.PDMat(jl(A), Cholesky(jl(Matrix(c.factors)), c.uplo, c.info))
@@ -36,14 +35,12 @@ end
             @test batched isa JLArray
             @test Array(batched) ≈ logpdf(host, xs) rtol = rtol
 
-            # the dispatch is on any GPU array, not just CUDA, so a second backend routes here
             @test NormalizingFlows._device_specific_logpdf(dev, xs_dev) ≈ batched rtol =
                 rtol
-            # and the host path is still Distributions
             @test NormalizingFlows._device_specific_logpdf(host, xs) == logpdf(host, xs)
 
-            # The gradient is the part that used to fail: a full covariance leaves one
-            # cotangent per differentiated use and adding them indexes the device array.
+            # A full covariance leaves one cotangent per differentiated use, and summing them
+            # indexes the device array.
             g = only(
                 Zygote.gradient(
                     x -> sum(NormalizingFlows._batched_mvnormal_logpdf(dev, x)), xs_dev
diff --git a/test/ext/CUDA/cuda.jl b/test/ext/CUDA/cuda.jl
index d1245f1f..10b3407f 100644
--- a/test/ext/CUDA/cuda.jl
+++ b/test/ext/CUDA/cuda.jl
@@ -28,15 +28,12 @@ function Bijectors._transform(flow::PlanarLayer, z::CuArray{T}) where {T<:Real}
     û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
     wT_z = Bijectors.aT_b(w, z)
 
-    # `flow.b` holds one element, so broadcasting it is the same as Bijectors' `first(flow.b)`
-    # without reading back from the device.
+    # `flow.b` holds one element, so broadcasting it avoids reading back from the device.
     tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, flow.b))
     transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
 
     return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
 end
-# Only the batched path is covered. A `CuVector` still falls to Bijectors' own method, whose
-# `first(flow.b)` reads back from the device, so single samples need `allowscalar(true)`.
 function Bijectors.with_logabsdet_jacobian(
     flow::PlanarLayer, z::CuMatrix{T}
 ) where {T<:Real}
@@ -115,9 +112,8 @@ end
     @test Array(x_back) ≈ x_cpu rtol = 1.0f-4
 end
 
-# `Distributions.logpdf` maps over columns and returns a host array, which used to break the
-# batched ELBO on a GPU. Planar layers are used throughout because coupling layers still
-# partition through a host sparse `PartitionMask`.
+# Planar layers throughout, because coupling layers partition through a host sparse
+# `PartitionMask`.
 @testset "batched ELBO on CUDA" begin
     CUDA.allowscalar(false)
     q0 = MvNormal(CUDA.zeros(Float32, 2), cu(Matrix{Float32}(I, 2, 2)))
@@ -144,8 +140,6 @@ end
         @test all(isfinite, Array(elbos))
         @test isfinite(elbo_batch(flow, logp, xs))
 
-        # The same flow on the host has to produce the same numbers, otherwise a wrong value
-        # in the planar overrides or the ELBO assembly would pass on types alone.
         cpu_q0 = MvNormal(zeros(Float32, 2), Matrix{Float32}(I, 2, 2))
         cpu_pl = fmap(Array, pl)
         cpu_flow = Bijectors.transformed(cpu_q0, cpu_pl)
@@ -154,9 +148,8 @@ end
         cpu_elbos = NormalizingFlows._batched_elbos(cpu_flow, cpu_logp, Array(xs))
         @test Array(elbos) ≈ cpu_elbos rtol = 1.0f-4
 
-        # A full covariance has to differentiate, not just evaluate: each differentiated use
-        # of one leaves a cotangent of a different matrix type, and summing them indexes the
-        # device array element by element.
+        # Each differentiated use of a full covariance leaves a cotangent of a different
+        # matrix type, and summing them indexes the device array.
         g = only(
             Zygote.gradient(x -> sum(NormalizingFlows._batched_elbos(flow, logp, x)), xs)
         )
@@ -170,8 +163,8 @@ end
     T = Float32
     d = 2
 
-    # Diagonal covariances exercise the broadcast branch of the log-density; the testset
-    # above covers the full covariance branch, gradient included.
+    # Diagonal covariances take the broadcast branch of the log-density. The full covariance
+    # branch is covered above.
     q0 = MvNormal(CUDA.zeros(T, d), Diagonal(CUDA.ones(T, d)))
     target = MvNormal(cu(T[2, -1]), Diagonal(CUDA.ones(T, d)))
     logp(z) = NormalizingFlows._device_specific_logpdf(target, z)
diff --git a/test/objectives.jl b/test/objectives.jl
index 8df79562..c5631803 100644
--- a/test/objectives.jl
+++ b/test/objectives.jl
@@ -37,8 +37,6 @@
 end
 
 @testset "batched MvNormal log-density" begin
-    # `_batched_elbos` needs a log-density that stays on the sample's device, so the whole
-    # array form has to agree with Distributions on every covariance type.
     @testset "$T" for T in (Float32, Float64)
         # a non-unit scale, so dropping the whitening is caught in this case too
         dists = (
@@ -48,16 +46,13 @@ end
         )
         @testset "$(nameof(typeof(d.Σ)))" for d in dists
             xs = randn(T, 3, 6)
-            # the gathered log-determinant has to agree with the one it replaces
             @test NormalizingFlows._cov_logdet(d.Σ) ≈ logdet(d.Σ) rtol = sqrt(eps(T))
 
             batched = NormalizingFlows._batched_mvnormal_logpdf(d, xs)
             @test eltype(batched) == T
             @test batched ≈ logpdf(d, xs) rtol = sqrt(eps(T))
-            # the generic fallback must leave the host path untouched
             @test NormalizingFlows._device_specific_logpdf(d, xs) == logpdf(d, xs)
 
-            # it sits inside the differentiated ELBO, so the pullback has to work too
             loss(v) = sum(NormalizingFlows._batched_mvnormal_logpdf(d, reshape(v, 3, 6)))
             g_ref = ForwardDiff.gradient(loss, vec(xs))
             @test all(isfinite, g_ref)

From 3bf79ad7c0051326263b7d64f9dd84841e9b3927 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 29 Aug 2026 22:02:47 +0530
Subject: [PATCH 17/17] Put one sentence per line in the docstrings

---
 example/README.md         | 3 ++-
 src/NormalizingFlows.jl   | 7 +++----
 src/flows/neuralspline.jl | 7 ++-----
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/example/README.md b/example/README.md
index 66045714..c4bc246f 100644
--- a/example/README.md
+++ b/example/README.md
@@ -17,4 +17,5 @@ using Pkg; Pkg.activate("."); Pkg.instantiate()
 ```
 This will install all needed packages, at the exact versions when the model was last updated. Then you can run the model code with `include("<example-to-run>.jl")`, or by running the example script line-by-line.
 
-`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA. Activate and instantiate `gpu/` to run it.
+`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA.
+Activate and instantiate `gpu/` to run it.
diff --git a/src/NormalizingFlows.jl b/src/NormalizingFlows.jl
index e327f827..f1a5a811 100644
--- a/src/NormalizingFlows.jl
+++ b/src/NormalizingFlows.jl
@@ -141,8 +141,8 @@ end
 """
     _device_draw(rng, s, dims)
 
-Draw a `dims`-shaped sample from `s` into an array on the device `rng` targets. Device
-extensions add the methods.
+Draw a `dims`-shaped sample from `s` into an array on the device `rng` targets.
+Device extensions add the methods.
 """
 function _device_draw end
 
@@ -154,8 +154,7 @@ function _device_draw end
     _device_specific_logpdf(d, xs)
 
 Log-density of `d` at each column of `xs`, left on the device holding `xs`.
-`Distributions.logpdf` maps over the columns and materialises a host array, so the ELBO
-cannot be assembled from it when the samples live on a GPU.
+`Distributions.logpdf` maps over the columns and materialises a host array, so the ELBO cannot be assembled from it when the samples live on a GPU.
 """
 _device_specific_logpdf(d, xs::AbstractMatrix) = logpdf(d, xs)
 
diff --git a/src/flows/neuralspline.jl b/src/flows/neuralspline.jl
index 6fad16c1..b207afc2 100644
--- a/src/flows/neuralspline.jl
+++ b/src/flows/neuralspline.jl
@@ -226,11 +226,8 @@ Returns
 - `Bijectors.TransformedDistribution` representing the NSF flow.
 
 !!! note
-    The rational quadratic spline is written with whole-array operations, so it runs on the
-    CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`,
-    `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`. The assembled flow is CPU only: the
-    coupling layer splits its input through `Bijectors.PartitionMask`, which holds host
-    sparse matrices.
+    The rational quadratic spline is written with whole-array operations, so it runs on the CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`, `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`.
+    The assembled flow is CPU only, because the coupling layer splits its input through `Bijectors.PartitionMask`, which holds host sparse matrices.
 
 !!! note
     When training the flow, mark the base distribution as a leaf first (`Functors.@leaf