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/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 diff --git a/src/bijectors/batched_rational_quadratic_spline.jl b/src/bijectors/batched_rational_quadratic_spline.jl new file mode 100644 index 00000000..db3eceb2 --- /dev/null +++ b/src/bijectors/batched_rational_quadratic_spline.jl @@ -0,0 +1,225 @@ +############################################# +### 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) + # 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 + +# 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) + # 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 + +""" + 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 + +# 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 = 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] +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 + + y = ifelse.(inside, y_bin, x) + 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 + +""" + 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. +""" +# 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{typeof(widths),typeof(heights),typeof(derivatives)}( + 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/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..ac2ba847 --- /dev/null +++ b/test/bijectors/batched_rational_quadratic_spline.jl @@ -0,0 +1,181 @@ +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 + @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 + +@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 + +@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 + +@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 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 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") 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