Batched, GPU-capable rational quadratic spline - #485
Draft
shravanngoswamii wants to merge 11 commits into
Draft
Conversation
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
|
Bijectors.jl documentation for PR #485 is available at: |
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…uild 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>
…he 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>
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>
…fferentiate 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>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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>
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.
Draft. Adds a batched, GPU-capable rational quadratic spline that works with any AD backend, no hand-written derivative rules. Goal is for NormalizingFlows nsf to use it instead of MonotonicSplines.
Built as plain array operations (reduction bin-search, differentiable gather, fused broadcast) so it runs on Array and CuArray from one code path and every backend differentiates it directly.
Progress:
CPU tests pass locally (1125 value checks over Float32/Float64, plus the AD suites). GPU tests run on CI.