Use the Bijectors batched RQS for nsf and drop MonotonicSplines#63
Draft
shravanngoswamii wants to merge 7 commits into
Draft
Use the Bijectors batched RQS for nsf and drop MonotonicSplines#63shravanngoswamii wants to merge 7 commits into
shravanngoswamii wants to merge 7 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>
7 tasks
# 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>
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.
nsf now builds its rational quadratic splines with the batched RQS in Bijectors (
rqs_params_from_raw,rqs_forward,rqs_inverse) instead of MonotonicSplines. That code is written with whole-array operations, so nsf trains under every supported AD backend and runs on the GPU. The Zygote-only restriction is gone, and MonotonicSplines is removed from the dependencies.The NSF ELBO AD test now covers Zygote, ReverseDiff, Enzyme, and Mooncake (Float32 and Float64), all passing locally.
Depends on TuringLang/Bijectors.jl#485. Until that is released, Bijectors is pinned to the branch via
[sources]inProject.toml. This resolves on Julia 1.11 and later; on 1.10[sources]is ignored, so this stays draft until the Bijectors release is tagged and the compat is pinned to it.