From 4245880597e3f4e548d3b8acafce80fbdeda1d2a Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sat, 18 Jul 2026 23:10:14 +0530
Subject: [PATCH 01/12] add batched RQS parameter constraints and raw-NN
 builder

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../batched_rational_quadratic_spline.jl      | 49 +++++++++++++++++++
 src/interface.jl                              |  1 +
 .../batched_rational_quadratic_spline.jl      | 32 ++++++++++++
 test/runtests.jl                              |  1 +
 4 files changed, 83 insertions(+)
 create mode 100644 src/bijectors/batched_rational_quadratic_spline.jl
 create mode 100644 test/bijectors/batched_rational_quadratic_spline.jl

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
new file mode 100644
index 00000000..028fa431
--- /dev/null
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -0,0 +1,49 @@
+#############################################
+### Batched rational quadratic spline (RQS) ###
+#############################################
+
+# A batched counterpart to `RationalQuadraticSpline` that evaluates many splines over a
+# batch of samples with whole-array operations, so the same source runs on `Array` and
+# `CuArray` and is differentiable by every AD backend without hand-written rules.
+#
+# Parameter arrays carry the knot axis first: `(K + 1, D, N)` for `K` bins, `D` transformed
+# dimensions, and `N` samples. Inputs are `(D, N)`.
+
+# Constrain raw parameters into a monotone knot grid on `[-B, B]`, batched along dim 1.
+# Mirrors the single-sample `RationalQuadraticSpline(widths, heights, derivatives, B)`
+# constructor: softmax to positive increments, cumulative sum to knots, scale to `[-B, B]`.
+function _rqs_constrain_knots(raw::AbstractArray, B)
+    T = eltype(raw)
+    Bc = T(B)
+    increments = LogExpFunctions.softmax(raw; dims=1)
+    lead = fill!(similar(raw, 1, Base.tail(size(raw))...), zero(T))
+    return cumsum(cat(lead, increments; dims=1); dims=1) .* (2 * Bc) .- Bc
+end
+
+# Interior derivatives are made positive with softplus; the endpoints are fixed to one so
+# the spline continues into the identity map outside `[-B, B]`.
+function _rqs_constrain_derivatives(raw::AbstractArray)
+    T = eltype(raw)
+    edge = fill!(similar(raw, 1, Base.tail(size(raw))...), one(T))
+    return cat(edge, LogExpFunctions.log1pexp.(raw), edge; dims=1)
+end
+
+"""
+    rqs_params_from_raw(θ_raw::AbstractMatrix, n_dims::Integer, B)
+
+Turn a matrix of raw neural-network outputs into constrained rational-quadratic-spline knot
+parameters. `θ_raw` has shape `((3K - 1) * n_dims, N)`, laid out per dimension as `K` width
+logits, `K` height logits, then `K - 1` derivative logits.
+
+Returns `(widths, heights, derivatives)`, each `(K + 1, n_dims, N)`, with `widths` and
+`heights` monotone on `[-B, B]` and `derivatives` positive with unit endpoints.
+"""
+function rqs_params_from_raw(θ_raw::AbstractMatrix, n_dims::Integer, B)
+    n_params, N = size(θ_raw)
+    K = (n_params ÷ n_dims + 1) ÷ 3
+    θ = reshape(θ_raw, 3K - 1, n_dims, N)
+    widths = _rqs_constrain_knots(θ[1:K, :, :], B)
+    heights = _rqs_constrain_knots(θ[(K + 1):(2K), :, :], B)
+    derivatives = _rqs_constrain_derivatives(θ[(2K + 1):(3K - 1), :, :])
+    return widths, heights, derivatives
+end
diff --git a/src/interface.jl b/src/interface.jl
index b0b3eb26..c149d1c9 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -390,6 +390,7 @@ include("bijectors/leaky_relu.jl")
 include("bijectors/coupling.jl")
 include("bijectors/normalise.jl")
 include("bijectors/rational_quadratic_spline.jl")
+include("bijectors/batched_rational_quadratic_spline.jl")
 
 ##################
 # Other includes #
diff --git a/test/bijectors/batched_rational_quadratic_spline.jl b/test/bijectors/batched_rational_quadratic_spline.jl
new file mode 100644
index 00000000..90180cf0
--- /dev/null
+++ b/test/bijectors/batched_rational_quadratic_spline.jl
@@ -0,0 +1,32 @@
+using Bijectors: rqs_params_from_raw
+
+@testset "batched RQS parameters" begin
+    @testset "T=$T, K=$K, D=$D, N=$N, B=$B" for T in (Float32, Float64),
+        K in (4, 8), D in (1, 3), N in (1, 16),
+        B in (2, 30)
+
+        θ_raw = randn(T, (3K - 1) * D, N)
+        widths, heights, derivatives = rqs_params_from_raw(θ_raw, D, B)
+
+        @test size(widths) == (K + 1, D, N)
+        @test size(heights) == (K + 1, D, N)
+        @test size(derivatives) == (K + 1, D, N)
+
+        # Raw parameters must not silently widen the element type.
+        @test eltype(widths) == T
+        @test eltype(heights) == T
+        @test eltype(derivatives) == T
+
+        for grid in (widths, heights)
+            # The first knot is exactly -B (a prepended zero before the cumsum); the last is
+            # B up to the floating-point error in the softmax normalisation.
+            @test all(grid[1, :, :] .== -T(B))
+            @test all(isapprox.(grid[end, :, :], T(B)))
+            @test all(diff(grid; dims=1) .> 0)
+        end
+
+        @test all(derivatives[1, :, :] .== one(T))
+        @test all(derivatives[end, :, :] .== one(T))
+        @test all(derivatives .> 0)
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 80f396dc..908f0d62 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -44,6 +44,7 @@ include("bijectors/utils.jl")
         include("normalising_flows.jl")
         include("bijectors/permute.jl")
         include("bijectors/rational_quadratic_spline.jl")
+        include("bijectors/batched_rational_quadratic_spline.jl")
         include("bijectors/named_bijector.jl")
         include("bijectors/leaky_relu.jl")
         include("bijectors/coupling.jl")

From 20c5d02118472ab2a75fcc6618cadb2f42b34c21 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 00:17:44 +0530
Subject: [PATCH 02/12] add batched RQS forward evaluation

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../batched_rational_quadratic_spline.jl      | 63 +++++++++++++++++++
 .../batched_rational_quadratic_spline.jl      | 51 ++++++++++++++-
 2 files changed, 113 insertions(+), 1 deletion(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index 028fa431..f0258aa6 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -47,3 +47,66 @@ function rqs_params_from_raw(θ_raw::AbstractMatrix, n_dims::Integer, B)
     derivatives = _rqs_constrain_derivatives(θ[(2K + 1):(3K - 1), :, :])
     return widths, heights, derivatives
 end
+
+# Locate the bin of each element and whether it lies inside the spline range. `knots` is
+# `(K + 1, D, N)`, `x` is `(D, N)`. `count` is the number of knots not exceeding `x`, so a
+# point in `[knots[k], knots[k+1])` gives `count == k`; `count == 0` or `count == K + 1`
+# means it is below or above the range. The comparison and integer reduction are
+# non-differentiable by construction, which is what confines the gradient to the arithmetic
+# of the selected bin.
+function _rqs_bin(knots::AbstractArray, x::AbstractMatrix)
+    K = size(knots, 1) - 1
+    count = dropdims(sum(knots .<= reshape(x, 1, size(x)...); dims=1); dims=1)
+    inside = (count .>= 1) .& (count .<= K)
+    return clamp.(count, 1, K), inside
+end
+
+# Gather the lower and upper knot values of each element's bin into `(D, N)` arrays. Linear
+# indices are built by broadcast so they live on the same device as the parameters, and the
+# gather is a plain `getindex` by integer array: vectorized on the GPU and differentiable on
+# every backend, with the gradient flowing back to the two selected knots.
+function _rqs_gather(knots::AbstractArray, k::AbstractMatrix{<:Integer})
+    stride1 = size(knots, 1)
+    D, N = size(k)
+    offset =
+        reshape(0:(D - 1), D, 1) .* stride1 .+ reshape(0:(N - 1), 1, N) .* (stride1 * D)
+    flat = reshape(knots, :)
+    return flat[offset .+ k], flat[offset .+ (k .+ 1)]
+end
+
+"""
+    rqs_forward(x, widths, heights, derivatives)
+
+Evaluate the batched rational-quadratic spline forward. `x` is `(D, N)` and the knot
+parameters are `(K + 1, D, N)`. Returns `(y, logjac)` with `y` of shape `(D, N)` and
+`logjac` of shape `(1, N)`, the per-sample sum over dimensions of `log|dy/dx|`. Outside
+`[widths[1], widths[end]]` the map is the identity and contributes zero to `logjac`.
+"""
+function rqs_forward(
+    x::AbstractMatrix,
+    widths::AbstractArray,
+    heights::AbstractArray,
+    derivatives::AbstractArray,
+)
+    T = eltype(x)
+    k, inside = _rqs_bin(widths, x)
+    xₖ, xₖ₊₁ = _rqs_gather(widths, k)
+    yₖ, yₖ₊₁ = _rqs_gather(heights, k)
+    dₖ, dₖ₊₁ = _rqs_gather(derivatives, k)
+
+    Δx = xₖ₊₁ .- xₖ
+    Δy = yₖ₊₁ .- yₖ
+    s = Δy ./ Δx
+    # Clamp keeps the discarded (out-of-range) branch finite so its zero-weighted gradient
+    # never turns into NaN; inside the range ξ is already in [0, 1] and clamp is a no-op.
+    ξ = clamp.((x .- xₖ) ./ Δx, zero(T), one(T))
+
+    denom = @. s + (dₖ₊₁ + dₖ - 2s) * ξ * (1 - ξ)
+    y_bin = @. yₖ + Δy * (s * ξ^2 + dₖ * ξ * (1 - ξ)) / denom
+    nom = @. dₖ₊₁ * ξ^2 + 2s * ξ * (1 - ξ) + dₖ * (1 - ξ)^2
+    logjac_bin = @. 2 * log(abs(s)) + log(abs(nom)) - 2 * log(abs(denom))
+
+    y = ifelse.(inside, y_bin, x)
+    logjac = ifelse.(inside, logjac_bin, zero(T))
+    return y, sum(logjac; dims=1)
+end
diff --git a/test/bijectors/batched_rational_quadratic_spline.jl b/test/bijectors/batched_rational_quadratic_spline.jl
index 90180cf0..4a25d144 100644
--- a/test/bijectors/batched_rational_quadratic_spline.jl
+++ b/test/bijectors/batched_rational_quadratic_spline.jl
@@ -1,4 +1,5 @@
-using Bijectors: rqs_params_from_raw
+using Bijectors: rqs_params_from_raw, rqs_forward, rqs_univariate
+using ForwardDiff: ForwardDiff
 
 @testset "batched RQS parameters" begin
     @testset "T=$T, K=$K, D=$D, N=$N, B=$B" for T in (Float32, Float64),
@@ -30,3 +31,51 @@ using Bijectors: rqs_params_from_raw
         @test all(derivatives .> 0)
     end
 end
+
+@testset "batched RQS forward" begin
+    @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
+        K in (4, 8), D in (1, 3),
+        N in (1, 8)
+
+        B = 5
+        w, h, d = rqs_params_from_raw(randn(T, (3K - 1) * D, N), D, B)
+        x = T(0.8B) .* (2 .* rand(T, D, N) .- 1)       # well inside [-B, B]
+        y, logjac = rqs_forward(x, w, h, d)
+
+        @test size(y) == (D, N)
+        @test size(logjac) == (1, N)
+        @test eltype(y) == T
+        @test eltype(logjac) == T
+
+        # Each column reproduces the legacy single-spline evaluation.
+        for n in 1:N, i in 1:D
+            @test y[i, n] ≈ rqs_univariate(w[:, i, n], h[:, i, n], d[:, i, n], x[i, n])
+        end
+
+        # logjac is the true log-derivative. The coupling is diagonal per dimension, so
+        # exp(logjac[n]) equals the product over dims of dyᵢ/dxᵢ from ForwardDiff.
+        for n in 1:N
+            prod_dydx = one(T)
+            for i in 1:D
+                scalar_forward = function (xi)
+                    xcol = reshape([j == i ? xi : x[j, n] for j in 1:D], D, 1)
+                    return rqs_forward(xcol, w[:, :, n:n], h[:, :, n:n], d[:, :, n:n])[1][i]
+                end
+                dydx = ForwardDiff.derivative(scalar_forward, x[i, n])
+                @test dydx > 0
+                prod_dydx *= dydx
+            end
+            @test prod_dydx ≈ exp(logjac[1, n])
+        end
+    end
+
+    @testset "out-of-range identity, T=$T" for T in (Float32, Float64)
+        B = 5
+        K, D, N = 6, 2, 4
+        w, h, d = rqs_params_from_raw(randn(T, (3K - 1) * D, N), D, B)
+        x = T[2B -2B 3B -3B; 2B -2B 3B -3B]           # all outside [-B, B]
+        y, logjac = rqs_forward(x, w, h, d)
+        @test y == x
+        @test all(iszero, logjac)
+    end
+end

From c31eefccd48de6b2b726204512f102921d173ffd Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 00:29:56 +0530
Subject: [PATCH 03/12] add batched RQS inverse evaluation

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../batched_rational_quadratic_spline.jl      | 52 +++++++++++++++++--
 .../batched_rational_quadratic_spline.jl      | 52 ++++++++++++++++++-
 2 files changed, 100 insertions(+), 4 deletions(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index f0258aa6..9bbd3780 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -103,10 +103,56 @@ function rqs_forward(
 
     denom = @. s + (dₖ₊₁ + dₖ - 2s) * ξ * (1 - ξ)
     y_bin = @. yₖ + Δy * (s * ξ^2 + dₖ * ξ * (1 - ξ)) / denom
-    nom = @. dₖ₊₁ * ξ^2 + 2s * ξ * (1 - ξ) + dₖ * (1 - ξ)^2
-    logjac_bin = @. 2 * log(abs(s)) + log(abs(nom)) - 2 * log(abs(denom))
 
     y = ifelse.(inside, y_bin, x)
-    logjac = ifelse.(inside, logjac_bin, zero(T))
+    logjac = ifelse.(inside, _rqs_forward_logjac(s, dₖ, dₖ₊₁, ξ), zero(T))
     return y, sum(logjac; dims=1)
 end
+
+# Forward log|dy/dx| for a bin at spline coordinate ξ in [0, 1], reused by the inverse.
+function _rqs_forward_logjac(s, dₖ, dₖ₊₁, ξ)
+    denom = @. s + (dₖ₊₁ + dₖ - 2s) * ξ * (1 - ξ)
+    nom = @. dₖ₊₁ * ξ^2 + 2s * ξ * (1 - ξ) + dₖ * (1 - ξ)^2
+    return @. 2 * log(abs(s)) + log(abs(nom)) - 2 * log(abs(denom))
+end
+
+"""
+    rqs_inverse(y, widths, heights, derivatives)
+
+Invert the batched rational-quadratic spline. `y` is `(D, N)` and the knot parameters are
+`(K + 1, D, N)`. Returns `(x, logjac)` with `x` of shape `(D, N)` and `logjac` of shape
+`(1, N)`, the per-sample sum over dimensions of `log|dx/dy|`. The bin holding each `y` is a
+monotone quadratic in the spline coordinate, solved with the roundoff-stable root; the
+log-det is the negation of the forward log-det at the recovered coordinate.
+"""
+function rqs_inverse(
+    y::AbstractMatrix,
+    widths::AbstractArray,
+    heights::AbstractArray,
+    derivatives::AbstractArray,
+)
+    T = eltype(y)
+    k, inside = _rqs_bin(heights, y)
+    xₖ, xₖ₊₁ = _rqs_gather(widths, k)
+    yₖ, yₖ₊₁ = _rqs_gather(heights, k)
+    dₖ, dₖ₊₁ = _rqs_gather(derivatives, k)
+
+    Δx = xₖ₊₁ .- xₖ
+    Δy = yₖ₊₁ .- yₖ
+    s = Δy ./ Δx
+    # Clamp to the bin so the discarded (out-of-range) branch stays finite; inside the range
+    # Δy2 already lies in [0, Δy] and clamp is a no-op.
+    Δy2 = clamp.(y .- yₖ, zero(T), Δy)
+
+    c1 = dₖ₊₁ .+ dₖ .- 2 .* s
+    a = @. Δy * (s - dₖ) + Δy2 * c1
+    b = @. Δy * dₖ - Δy2 * c1
+    c = @. -s * Δy2
+    disc = @. max(b^2 - 4 * a * c, zero(T))
+    denom = @. -b - sqrt(disc)
+    ξ = clamp.((2 .* c) ./ denom, zero(T), one(T))
+
+    x = ifelse.(inside, xₖ .+ ξ .* Δx, y)
+    logjac = ifelse.(inside, .-_rqs_forward_logjac(s, dₖ, dₖ₊₁, ξ), zero(T))
+    return x, sum(logjac; dims=1)
+end
diff --git a/test/bijectors/batched_rational_quadratic_spline.jl b/test/bijectors/batched_rational_quadratic_spline.jl
index 4a25d144..eeda0dbd 100644
--- a/test/bijectors/batched_rational_quadratic_spline.jl
+++ b/test/bijectors/batched_rational_quadratic_spline.jl
@@ -1,4 +1,4 @@
-using Bijectors: rqs_params_from_raw, rqs_forward, rqs_univariate
+using Bijectors: rqs_params_from_raw, rqs_forward, rqs_inverse, rqs_univariate
 using ForwardDiff: ForwardDiff
 
 @testset "batched RQS parameters" begin
@@ -79,3 +79,53 @@ end
         @test all(iszero, logjac)
     end
 end
+
+@testset "batched RQS inverse" begin
+    @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
+        K in (4, 8), D in (1, 3),
+        N in (1, 8)
+
+        B = 5
+        w, h, d = rqs_params_from_raw(randn(T, (3K - 1) * D, N), D, B)
+        rtol = T == Float32 ? 1.0f-4 : 1.0e-9
+
+        # inverse ∘ forward
+        x = T(0.8B) .* (2 .* rand(T, D, N) .- 1)
+        y, logjac_fwd = rqs_forward(x, w, h, d)
+        xback, logjac_inv = rqs_inverse(y, w, h, d)
+        @test xback ≈ x rtol = rtol
+        @test logjac_inv ≈ -logjac_fwd rtol = rtol
+
+        # forward ∘ inverse (heights are also constrained to [-B, B])
+        yin = T(0.8B) .* (2 .* rand(T, D, N) .- 1)
+        xr, _ = rqs_inverse(yin, w, h, d)
+        yr, _ = rqs_forward(xr, w, h, d)
+        @test yr ≈ yin rtol = rtol
+    end
+
+    @testset "out-of-range identity, T=$T" for T in (Float32, Float64)
+        B = 5
+        K, D, N = 6, 2, 4
+        w, h, d = rqs_params_from_raw(randn(T, (3K - 1) * D, N), D, B)
+        y = T[2B -2B 3B -3B; 2B -2B 3B -3B]
+        x, logjac = rqs_inverse(y, w, h, d)
+        @test x == y
+        @test all(iszero, logjac)
+    end
+
+    @testset "boundary gradient is finite, T=$T" for T in (Float32, Float64)
+        B = 5
+        K, D, N = 6, 1, 5
+        w, h, d = rqs_params_from_raw(randn(T, (3K - 1) * D, N), D, B)
+        # points spanning below, on, and above the range
+        y = reshape(T[-B - 1, -B, 0, B, B + 1], D, N)
+        g = ForwardDiff.gradient(
+            v -> sum(rqs_inverse(reshape(v, D, N), w, h, d)[1]), vec(y)
+        )
+        @test all(isfinite, g)
+        gj = ForwardDiff.gradient(
+            v -> sum(rqs_inverse(reshape(v, D, N), w, h, d)[2]), vec(y)
+        )
+        @test all(isfinite, gj)
+    end
+end

From ce926eee88679a32e5217ae3c8b9a9547ec55d48 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 01:42:10 +0530
Subject: [PATCH 04/12] add BatchedRQS bijector interface

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../batched_rational_quadratic_spline.jl      | 53 +++++++++++++++++++
 .../batched_rational_quadratic_spline.jl      | 47 +++++++++++++++-
 2 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index 9bbd3780..2859e1d3 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -156,3 +156,56 @@ function rqs_inverse(
     logjac = ifelse.(inside, .-_rqs_forward_logjac(s, dₖ, dₖ₊₁, ξ), zero(T))
     return x, sum(logjac; dims=1)
 end
+
+"""
+    BatchedRQS(widths, heights, derivatives)
+    BatchedRQS(θ_raw::AbstractMatrix, n_dims::Integer, B)
+
+Batched rational quadratic spline bijector. The knot parameters `widths`, `heights`, and
+`derivatives` are `(K + 1, D, N)` arrays already constrained to a valid monotone spline (as
+produced by [`rqs_params_from_raw`](@ref)); the second constructor builds them from raw
+neural-network outputs. `transform` maps `(D, N)` inputs to `(D, N)` outputs and
+`logabsdetjac` returns the per-sample `(N,)` log-determinant.
+"""
+struct BatchedRQS{T<:AbstractArray} <: Bijector
+    widths::T
+    heights::T
+    derivatives::T
+
+    function BatchedRQS(widths::T, heights::T, derivatives::T) where {T<:AbstractArray}
+        if !(size(widths) == size(heights) == size(derivatives))
+            throw(
+                DimensionMismatch("widths, heights, and derivatives must share their shape")
+            )
+        end
+        return new{T}(widths, heights, derivatives)
+    end
+end
+
+function BatchedRQS(θ_raw::AbstractMatrix, n_dims::Integer, B)
+    return BatchedRQS(rqs_params_from_raw(θ_raw, n_dims, B)...)
+end
+
+function transform(b::BatchedRQS, x::AbstractMatrix)
+    return first(rqs_forward(x, b.widths, b.heights, b.derivatives))
+end
+
+function transform(ib::Inverse{<:BatchedRQS}, y::AbstractMatrix)
+    b = ib.orig
+    return first(rqs_inverse(y, b.widths, b.heights, b.derivatives))
+end
+
+function logabsdetjac(b::BatchedRQS, x::AbstractMatrix)
+    return vec(last(rqs_forward(x, b.widths, b.heights, b.derivatives)))
+end
+
+function with_logabsdet_jacobian(b::BatchedRQS, x::AbstractMatrix)
+    y, logjac = rqs_forward(x, b.widths, b.heights, b.derivatives)
+    return y, vec(logjac)
+end
+
+function with_logabsdet_jacobian(ib::Inverse{<:BatchedRQS}, y::AbstractMatrix)
+    b = ib.orig
+    x, logjac = rqs_inverse(y, b.widths, b.heights, b.derivatives)
+    return x, vec(logjac)
+end
diff --git a/test/bijectors/batched_rational_quadratic_spline.jl b/test/bijectors/batched_rational_quadratic_spline.jl
index eeda0dbd..dd693d1e 100644
--- a/test/bijectors/batched_rational_quadratic_spline.jl
+++ b/test/bijectors/batched_rational_quadratic_spline.jl
@@ -1,4 +1,5 @@
-using Bijectors: rqs_params_from_raw, rqs_forward, rqs_inverse, rqs_univariate
+using Bijectors: rqs_params_from_raw, rqs_forward, rqs_inverse, rqs_univariate, BatchedRQS
+using Bijectors: transform, with_logabsdet_jacobian, logabsdetjac, inverse
 using ForwardDiff: ForwardDiff
 
 @testset "batched RQS parameters" begin
@@ -129,3 +130,47 @@ end
         @test all(isfinite, gj)
     end
 end
+
+@testset "BatchedRQS bijector" begin
+    @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
+        K in (4, 8), D in (1, 3),
+        N in (1, 8)
+
+        B = 5
+        θ_raw = randn(T, (3K - 1) * D, N)
+        w, h, d = rqs_params_from_raw(θ_raw, D, B)
+        b = BatchedRQS(w, h, d)
+        x = T(0.8B) .* (2 .* rand(T, D, N) .- 1)
+        rtol = T == Float32 ? 1.0f-4 : 1.0e-9
+
+        # The convenience constructor matches building from constrained params.
+        b2 = BatchedRQS(θ_raw, D, B)
+        @test b2.widths == w && b2.heights == h && b2.derivatives == d
+
+        # transform / logabsdetjac / with_logabsdet_jacobian are consistent.
+        y = transform(b, x)
+        ladj = logabsdetjac(b, x)
+        y2, ladj2 = with_logabsdet_jacobian(b, x)
+        @test y == y2
+        @test ladj == ladj2
+        @test size(y) == (D, N)
+        @test size(ladj) == (N,)
+        @test eltype(y) == T
+        @test eltype(ladj) == T
+
+        # Inverse via the generic wrapper round-trips and negates the log-det.
+        xback, ladj_inv = with_logabsdet_jacobian(inverse(b), y)
+        @test xback ≈ x rtol = rtol
+        @test ladj_inv ≈ -ladj rtol = rtol
+        @test transform(inverse(b), y) ≈ x rtol = rtol
+
+        # Type stability.
+        @inferred transform(b, x)
+        @inferred with_logabsdet_jacobian(b, x)
+    end
+
+    @testset "shape mismatch is rejected" begin
+        w = randn(Float64, 5, 2, 3)
+        @test_throws DimensionMismatch BatchedRQS(w, w, randn(Float64, 5, 2, 4))
+    end
+end

From 102b3bcf8c1169380027e6abb4091dfc5af313f2 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 01:49:39 +0530
Subject: [PATCH 05/12] cover batched RQS under the AD integration suites

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 test/test_resources.jl | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/test/test_resources.jl b/test/test_resources.jl
index ed71100b..f61f8561 100644
--- a/test/test_resources.jl
+++ b/test/test_resources.jl
@@ -67,7 +67,12 @@ ADTestCase(name::String, func, arg) = ADTestCase(name, func, arg, :_default)
 _settag(c::ADTestCase, tag::Symbol) = ADTestCase(c.name, c.func, c.arg, tag)
 
 const _AD_TAGS = (
-    :veccorrbijector, :veccholeskybijector, :planarlayer, :pdvecbijector, :stackedbijector
+    :veccorrbijector,
+    :veccholeskybijector,
+    :planarlayer,
+    :pdvecbijector,
+    :stackedbijector,
+    :batchedrqs,
 )
 
 """
@@ -171,6 +176,33 @@ function _gen_testcases(::Val{:veccholeskybijector})
     return cases
 end
 
+function _gen_testcases(::Val{:batchedrqs})
+    rng = _testcase_rng()
+    K, D, N, B = 4, 2, 3, 5
+    n_raw = (3K - 1) * D
+
+    # θ packs the raw per-sample spline parameters followed by the batched inputs, so the
+    # gradient covers the full path: constraints, gather, spline, and log-det.
+    forward = function (θ)
+        raw = reshape(θ[1:(n_raw * N)], n_raw, N)
+        x = reshape(θ[(n_raw * N + 1):end], D, N)
+        b = Bijectors.BatchedRQS(raw, D, B)
+        return sum(transform(b, x)) + sum(logabsdetjac(b, x))
+    end
+    backward = function (θ)
+        raw = reshape(θ[1:(n_raw * N)], n_raw, N)
+        y = reshape(θ[(n_raw * N + 1):end], D, N)
+        b = Bijectors.BatchedRQS(raw, D, B)
+        binv = inverse(b)
+        return sum(transform(binv, y)) + sum(logabsdetjac(binv, y))
+    end
+    arg = randn(rng, n_raw * N + D * N)
+    return [
+        ADTestCase("BatchedRQS forward", forward, arg),
+        ADTestCase("BatchedRQS inverse", backward, copy(arg)),
+    ]
+end
+
 function _gen_testcases(::Val{:planarlayer})
     rng = _testcase_rng()
     # logpdf of a flow with a planar layer and two-dimensional inputs

From da51dab8a5c53002ced34987f42bf4e0010055a0 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 01:58:04 +0530
Subject: [PATCH 06/12] let BatchedRQS fields hold different array types so AD
 backends can build it

ReverseDiff returns the constrained parameter arrays as different tracked types, which
a single shared type parameter rejected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .../batched_rational_quadratic_spline.jl      | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index 2859e1d3..bec68409 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -167,18 +167,25 @@ produced by [`rqs_params_from_raw`](@ref)); the second constructor builds them f
 neural-network outputs. `transform` maps `(D, N)` inputs to `(D, N)` outputs and
 `logabsdetjac` returns the per-sample `(N,)` log-determinant.
 """
-struct BatchedRQS{T<:AbstractArray} <: Bijector
-    widths::T
-    heights::T
-    derivatives::T
-
-    function BatchedRQS(widths::T, heights::T, derivatives::T) where {T<:AbstractArray}
+# The three fields carry independent type parameters because automatic differentiation can
+# return them as different array types (for example ReverseDiff yields a TrackedArray for one
+# and an Array of tracked scalars for another).
+struct BatchedRQS{Tw<:AbstractArray,Th<:AbstractArray,Td<:AbstractArray} <: Bijector
+    widths::Tw
+    heights::Th
+    derivatives::Td
+
+    function BatchedRQS(
+        widths::AbstractArray, heights::AbstractArray, derivatives::AbstractArray
+    )
         if !(size(widths) == size(heights) == size(derivatives))
             throw(
                 DimensionMismatch("widths, heights, and derivatives must share their shape")
             )
         end
-        return new{T}(widths, heights, derivatives)
+        return new{typeof(widths),typeof(heights),typeof(derivatives)}(
+            widths, heights, derivatives
+        )
     end
 end
 

From da8c002be6da70694a4200a2e193ee9ba47ed767 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 02:17:23 +0530
Subject: [PATCH 07/12] keep the batched RQS gather index on-device so it runs
 unchanged on the GPU

The linear index is now fused into one broadcast with the bin array, instead of adding a
host offset array to a device index, which would fail under CUDA.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 src/bijectors/batched_rational_quadratic_spline.jl | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index bec68409..bdef8fcd 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -68,10 +68,13 @@ end
 function _rqs_gather(knots::AbstractArray, k::AbstractMatrix{<:Integer})
     stride1 = size(knots, 1)
     D, N = size(k)
-    offset =
-        reshape(0:(D - 1), D, 1) .* stride1 .+ reshape(0:(N - 1), 1, N) .* (stride1 * D)
+    di = reshape(1:D, D, 1)
+    ni = reshape(1:N, 1, N)
+    # Fuse the linear index into a single broadcast that includes `k`, so the result lives on
+    # the same device as `k`; a standalone host offset array added to a device `k` would fail.
+    lin = @. k + (di - 1) * stride1 + (ni - 1) * (stride1 * D)
     flat = reshape(knots, :)
-    return flat[offset .+ k], flat[offset .+ (k .+ 1)]
+    return flat[lin], flat[lin .+ 1]
 end
 
 """

From e61ce3155fc48a37d3a4a588bac6748b65132e9f Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 02:17:23 +0530
Subject: [PATCH 08/12] add a CUDA test suite for the batched RQS and run it on
 buildkite

Checks device execution without scalar indexing, host/device agreement, and that the
Zygote gradient on the GPU matches the CPU.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 .buildkite/pipeline.yml |  7 ++++-
 test/gpu/Project.toml   | 13 +++++++++
 test/gpu/main.jl        | 60 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 test/gpu/Project.toml
 create mode 100644 test/gpu/main.jl

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 4803bbff..fc4773f6 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -10,7 +10,12 @@ steps:
       #     dirs:
       #       - src
       #       - ext
-    command: julia --eval='println("Skipping CUDA tests - pipeline configured to do nothing")'
+    # `Pkg.develop` makes Julia 1.10 use the working tree's Bijectors. The `[sources]` section
+    # in test/gpu/Project.toml only takes effect on Julia 1.11+, so on 1.10 the registry
+    # version would be resolved instead.
+    command: |
+      julia --project=test/gpu --color=yes -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+      julia --project=test/gpu --color=yes test/gpu/main.jl
     agents:
       queue: "cuda"
     if: build.message !~ /\[skip tests\]/
diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml
new file mode 100644
index 00000000..93fbbb11
--- /dev/null
+++ b/test/gpu/Project.toml
@@ -0,0 +1,13 @@
+[deps]
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[sources]
+Bijectors = {path = "../.."}
+
+[compat]
+CUDA = "5"
+Zygote = "0.6, 0.7"
diff --git a/test/gpu/main.jl b/test/gpu/main.jl
new file mode 100644
index 00000000..b465ddb8
--- /dev/null
+++ b/test/gpu/main.jl
@@ -0,0 +1,60 @@
+using Bijectors
+using Bijectors: BatchedRQS, transform, logabsdetjac, with_logabsdet_jacobian, inverse
+using CUDA
+using Random
+using Test
+using Zygote
+
+# Scalar indexing on a GPU array is the usual sign that an operation fell back to a slow,
+# element-by-element host loop. Disallowing it turns any such fallback into a test failure.
+CUDA.allowscalar(false)
+
+@testset "Batched RQS on CUDA" begin
+    if !CUDA.functional()
+        @info "CUDA is not functional on this agent, skipping GPU tests"
+    else
+        rng = MersenneTwister(1)
+        K, D, N = 4, 3, 8
+        B = 5.0f0
+        n_raw = (3K - 1) * D
+
+        θ_cpu = randn(rng, Float32, n_raw, N)
+        x_cpu = randn(rng, Float32, D, N)
+        θ_gpu = cu(θ_cpu)
+        x_gpu = cu(x_cpu)
+
+        @testset "runs on device and keeps Float32" begin
+            b = BatchedRQS(θ_gpu, D, B)
+            y = transform(b, x_gpu)
+            lad = logabsdetjac(b, x_gpu)
+            x_back = transform(inverse(b), y)
+            @test y isa CuArray{Float32}
+            @test lad isa CuArray{Float32}
+            @test x_back isa CuArray{Float32}
+            @test size(y) == (D, N)
+            @test size(lad) == (N,)
+        end
+
+        @testset "device result matches host" begin
+            b_cpu = BatchedRQS(θ_cpu, D, B)
+            b_gpu = BatchedRQS(θ_gpu, D, B)
+            y_cpu, lad_cpu = with_logabsdet_jacobian(b_cpu, x_cpu)
+            y_gpu, lad_gpu = with_logabsdet_jacobian(b_gpu, x_gpu)
+            @test Array(y_gpu) ≈ y_cpu rtol = 1.0f-4
+            @test Array(lad_gpu) ≈ lad_cpu rtol = 1.0f-4
+
+            xb_cpu = transform(inverse(b_cpu), y_cpu)
+            xb_gpu = transform(inverse(b_gpu), y_gpu)
+            @test Array(xb_gpu) ≈ xb_cpu rtol = 1.0f-4
+            @test Array(xb_gpu) ≈ x_cpu rtol = 1.0f-4
+        end
+
+        @testset "gradient on device matches host" begin
+            loss(θ, x) = sum(logabsdetjac(BatchedRQS(θ, D, B), x))
+            g_cpu = only(Zygote.gradient(θ -> loss(θ, x_cpu), θ_cpu))
+            g_gpu = only(Zygote.gradient(θ -> loss(θ, x_gpu), θ_gpu))
+            @test g_gpu isa CuArray{Float32}
+            @test Array(g_gpu) ≈ g_cpu rtol = 1.0f-3
+        end
+    end
+end

From 229ad56c569d311ffb800b27709ad14bac0218de Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 03:17:07 +0530
Subject: [PATCH 09/12] build the batched RQS boundary rows without mutation so
 Zygote can differentiate it

The unit and zero endpoint rows were created with fill!, which Zygote rejects. They are now
built by broadcasting, keeping the array type and every AD backend working.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 src/bijectors/batched_rational_quadratic_spline.jl | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index bdef8fcd..987e17f0 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -16,7 +16,9 @@ function _rqs_constrain_knots(raw::AbstractArray, B)
     T = eltype(raw)
     Bc = T(B)
     increments = LogExpFunctions.softmax(raw; dims=1)
-    lead = fill!(similar(raw, 1, Base.tail(size(raw))...), zero(T))
+    # A leading zero row, built without mutation so Zygote can differentiate it, and from a
+    # slice of `increments` so it keeps the array type (`Array`, `CuArray`, ...).
+    lead = zero(T) .* increments[1:1, :, :]
     return cumsum(cat(lead, increments; dims=1); dims=1) .* (2 * Bc) .- Bc
 end
 
@@ -24,7 +26,8 @@ end
 # the spline continues into the identity map outside `[-B, B]`.
 function _rqs_constrain_derivatives(raw::AbstractArray)
     T = eltype(raw)
-    edge = fill!(similar(raw, 1, Base.tail(size(raw))...), one(T))
+    # Unit endpoint rows, built without mutation (see `_rqs_constrain_knots`).
+    edge = zero(T) .* raw[1:1, :, :] .+ one(T)
     return cat(edge, LogExpFunctions.log1pexp.(raw), edge; dims=1)
 end
 

From 3ed048e7aaa5bdb02d37b71d6191dd570a396065 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 03:20:46 +0530
Subject: [PATCH 10/12] format the batched RQS tests with the blue style

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 test/bijectors/batched_rational_quadratic_spline.jl | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/test/bijectors/batched_rational_quadratic_spline.jl b/test/bijectors/batched_rational_quadratic_spline.jl
index dd693d1e..ac2ba847 100644
--- a/test/bijectors/batched_rational_quadratic_spline.jl
+++ b/test/bijectors/batched_rational_quadratic_spline.jl
@@ -4,7 +4,9 @@ using ForwardDiff: ForwardDiff
 
 @testset "batched RQS parameters" begin
     @testset "T=$T, K=$K, D=$D, N=$N, B=$B" for T in (Float32, Float64),
-        K in (4, 8), D in (1, 3), N in (1, 16),
+        K in (4, 8),
+        D in (1, 3),
+        N in (1, 16),
         B in (2, 30)
 
         θ_raw = randn(T, (3K - 1) * D, N)
@@ -35,7 +37,8 @@ end
 
 @testset "batched RQS forward" begin
     @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
-        K in (4, 8), D in (1, 3),
+        K in (4, 8),
+        D in (1, 3),
         N in (1, 8)
 
         B = 5
@@ -83,7 +86,8 @@ end
 
 @testset "batched RQS inverse" begin
     @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
-        K in (4, 8), D in (1, 3),
+        K in (4, 8),
+        D in (1, 3),
         N in (1, 8)
 
         B = 5
@@ -133,7 +137,8 @@ end
 
 @testset "BatchedRQS bijector" begin
     @testset "T=$T, K=$K, D=$D, N=$N" for T in (Float32, Float64),
-        K in (4, 8), D in (1, 3),
+        K in (4, 8),
+        D in (1, 3),
         N in (1, 8)
 
         B = 5

From 2670fd7db9b8a8a22d469eb1e24ec4945acc647d Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Sun, 19 Jul 2026 03:51:55 +0530
Subject: [PATCH 11/12] derive the RQS gather indices from the bin array so
 they run on the GPU

Row and column indices came from reshaped host ranges, which cannot broadcast against a GPU
array. They are now built from the bin array with cumsum, keeping the whole gather on device.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 src/bijectors/batched_rational_quadratic_spline.jl | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl
index 987e17f0..db3eceb2 100644
--- a/src/bijectors/batched_rational_quadratic_spline.jl
+++ b/src/bijectors/batched_rational_quadratic_spline.jl
@@ -70,11 +70,12 @@ end
 # every backend, with the gradient flowing back to the two selected knots.
 function _rqs_gather(knots::AbstractArray, k::AbstractMatrix{<:Integer})
     stride1 = size(knots, 1)
-    D, N = size(k)
-    di = reshape(1:D, D, 1)
-    ni = reshape(1:N, 1, N)
-    # Fuse the linear index into a single broadcast that includes `k`, so the result lives on
-    # the same device as `k`; a standalone host offset array added to a device `k` would fail.
+    D = size(k, 1)
+    # Row and column indices of each entry, derived from `k` so they share its array type: a
+    # reshaped host range cannot take part in a broadcast against a GPU array.
+    unit = one.(k)
+    di = cumsum(unit; dims=1)
+    ni = cumsum(unit; dims=2)
     lin = @. k + (di - 1) * stride1 + (ni - 1) * (stride1 * D)
     flat = reshape(knots, :)
     return flat[lin], flat[lin .+ 1]

From ec8b05c726026d9204ba2eeb4cdfc07702119ed1 Mon Sep 17 00:00:00 2001
From: Shravan Goswami <shravanngoswamii@gmail.com>
Date: Wed, 5 Aug 2026 16:37:19 +0530
Subject: [PATCH 12/12] format

---
 docs/src/vector.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/src/vector.md b/docs/src/vector.md
index 6c4e62ac..c5188842 100644
--- a/docs/src/vector.md
+++ b/docs/src/vector.md
@@ -7,6 +7,7 @@ It assumes that there are three forms of samples from a distribution `d` that we
  1. **The original form**, which is what `rand(d)` returns.
 
  2. **A vectorised form**, which is a vector that contains a flattened version of the original form.
+
  3. **A linked vectorised form**, which is a vector in which:
 
       + each element is independent; and