The astronomical GP without the padded grid - #239
Draft
chrisfinlay wants to merge 2 commits into
Draft
chrisfinlay wants to merge 2 commits into
chrisfinlay wants to merge 2 commits into
Conversation
chrisfinlay
force-pushed
the
dft-gp-ast
branch
from
September 13, 2026 01:48
3424d43 to
8b1934d
Compare
chrisfinlay
changed the base branch from
poly-interp-reference
to
poly-interp-hybrid
September 13, 2026 01:48
chrisfinlay
force-pushed
the
dft-gp-ast
branch
from
September 13, 2026 05:51
8b1934d to
278829c
Compare
chrisfinlay
changed the title
chrisfinlay
force-pushed
the
setup-cost
branch
from
September 13, 2026 13:21
b8fc890 to
0e64f8f
Compare
chrisfinlay
force-pushed
the
dft-gp-ast
branch
from
September 13, 2026 13:21
f85f480 to
8e56769
Compare
chrisfinlay
changed the title
…added grid The Fourier GPs reach a signal by padding the surviving modes up to a grid twice the data axis in each of frequency and time, inverse-transforming the whole of it and cropping the padding away. The padded grid is never wanted for itself: in two dimensions it is about four times the elements of the visibilities, formed and discarded once per block in the forward pass and again in reverse, and the astronomical prior keeps very few modes -- at the scaling benchmark's settings, 1 by 55 of a 16 by 300 grid. That chain is linear and separable, so the cropped output is a discrete Fourier transform of the surviving modes evaluated at the positions that survive the crop: one small matrix per axis, contracted in whichever order makes the intermediate smaller. `latent_to_signal_dft` is that, and `ast_vis:GPVisAstDFT` is `GPVisAst` with only this step replaced, so the prior, the initialisation and the blocking are the parent's and the two can be compared directly. The phases are reduced as exact integer ratios, so the matrices are the same in single precision as in double. The scan over baselines exists to bound the padded grid. With none to bound, and with its own stack of per-step outputs to pay for, it is a cost on this transform, so `auto` is a single step here. Measured on a 2016-baseline grid at 150 integrations, value and gradient together, against the FFT component at whichever block suited it best: 273.9 MB and 36.5 ms become 190.6 MB and 3.4 ms at 8 channels, 828.6 and 153.2 become 683.8 and 25.2 at 32, and 2386.4 and 596.0 become 1799.1 and 86.1 at 128, where the FFT route runs out of memory on a 6 GB card if it is not blocked. Tests hold the transform to the chain it stands for over five grid shapes, including a single channel, non-default padding and a supersampled axis, and hold the component to `GPVisAst` in value and in gradient.
chrisfinlay
force-pushed
the
dft-gp-ast
branch
from
September 13, 2026 14:39
8e56769 to
a644938
Compare
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.
Stacked on #242, which is the bottom of this stack and branches off
main.Nothing here touches the compiled kernels.
The Fourier GPs reach a signal by padding the surviving modes to a grid twice
the data axis in each of frequency and time, inverse-transforming the whole of
it and cropping the padding away. The chain is linear and separable, so the
cropped output is a discrete Fourier transform of the surviving modes evaluated
at the positions that survive the crop — one small matrix product per axis
instead of a grid about four times the size of the visibilities, formed and
discarded once per block in the forward pass and again in reverse.
ast_vis:GPVisAstDFTisGPVisAstwith that substitution made:The prior, the initialisation and the blocking over baselines are the parent's;
only the step from the surviving modes to the visibilities differs. The modes of
this component are set by the array's geometry rather than by a choice, so the
matrices are fixed at setup and shared by every baseline.
Measured
One GH200, 2016 baselines at 150 integrations, value and gradient together, each
case in its own process — the peak is a process high-water mark, not an
allocator figure.
GPVisAst, best blockGPVisAst, unblockedGPVisAstDFTThe two components agree to round-off in value and in gradient, which is what
the tests check — against the parent's own pads and crop, so the comparison is
to the code being replaced rather than to a second implementation of the same
idea.
autoon this component is a single step over every baseline. Its onlytransient is the intermediate of two matrix products, so the scan has nothing
large left to bound and its stack of per-step outputs would be a cost of its
own. An explicit
baseline_block_sizeis still honoured on both components.Scope
GPVisAstremains the default and is unchanged. This is a component you list inits place, so nothing moves unless a config asks for it.