Use the Bijectors batched RQS for nsf and drop MonotonicSplines - #63
Open
shravanngoswamii wants to merge 23 commits into
Open
Use the Bijectors batched RQS for nsf and drop MonotonicSplines#63shravanngoswamii wants to merge 23 commits into
shravanngoswamii wants to merge 23 commits into
Conversation
The neural spline coupling now calls Bijectors.rqs_params_from_raw, rqs_forward, and rqs_inverse, which are written with whole-array operations. nsf now trains under every supported AD backend and runs on the GPU, so the Zygote-only restriction is gone. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # Project.toml # test/Project.toml
Contributor
|
NormalizingFlows.jl documentation for PR #63 is available at: |
The neural spline demo needs the batched RQS, which is not yet released, so the example environment needs the same source pin as the main project. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…elow 0.13.189 Julia 1.10 ignores the [sources] pin, so the batched RQS was missing there. Enzyme 0.13.189 fails with an LLVM canonicalization error on 1.10; 0.13.188 is fine. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…constructor @leaf MvNormal now runs once in runtests.jl instead of inside every testset loop, and the base distributions use an identity covariance instead of the deprecated vector-of-sigmas form. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Enzyme 0.13.188 fails LLVM function verification when differentiating the batched RQS on Julia 1.10 (Float64 only); the other backends pass there and Enzyme passes on 1.11+. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…mal constructor MvNormal(mu, I) builds an IsoNormal, and Enzyme crashes differentiating destructure over that type. Diagonal(ones(T, 2)) keeps the exact DiagNormal the deprecated form produced. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
shravanngoswamii
marked this pull request as ready for review
August 19, 2026 10:25
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
nsfnow builds its rational quadratic splines with a batched implementation in this package (src/flows/rqs.jl) instead of MonotonicSplines, which is dropped as a dependency. The spline is written with whole-array operations, so it runs on the GPU and every AD backend differentiates it without hand-written rules. A small ReverseDiff extension strips tracking from the bin search and the root sign mask, which are piecewise constant and carry no gradient.The implementation was first proposed for Bijectors (TuringLang/Bijectors.jl#485) and moved here instead: it assumes a
(D, N)batch layout, which Bijectors deliberately does not commit to (#214, #370), while the coupling layers here already use that convention.Behaviour change: MonotonicSplines ignored
Bwhen scaling the knots and always scaled into (-5, 5]. Flows built with anyBother than 5 now define a different transform. Whether thensfdefault ofB = 30withK = 10is still sensible is worth a maintainer decision.The parameterisation follows the reference implementations: a minimum bin fraction and minimum derivative of 1e-3, boundary knots pinned exactly to
-BandB, and the cancellation-free form of the quadratic root for either sign of b.Tests:
test/rqs.jlchecks shapes, the parameter layout, monotone knots, the identity tails, round trips, agreement with the single-sample spline in Bijectors, and gradients against ForwardDiff under Zygote, ReverseDiff, Mooncake (reverse and forward), and Enzyme (Julia 1.11+, where it does not hit its LLVM verification bug), over interior inputs, mixed in-range and out-of-range batches, and parameters extreme enough to saturate the floors.test/ext/CUDA/cuda.jlchecks device execution without scalar indexing, host agreement, and the Zygote gradient on device.