fix: drop deleted TransformerNUFFTPyNUFFT from jax_profiling dataset setup - #129
Merged
Merged
Conversation
…setup
The transformer lookup dict in simulate() is built eagerly, so the reference
to al.TransformerNUFFTPyNUFFT (deleted by PyAutoArray#475) raised
AttributeError for EVERY instrument, not just the alma_high_res config that
selected it -- reproduced via simulate('sma'), a DFT dataset.
Drops the 'nufft_pynufft' arm and repoints alma_high_res at the nufftax-backed
TransformerNUFFT rather than DFT: its ~20GB dense-matrix OOM constraint is
real (5000 vis x 512x512 = 1.31e9, far above the ~1e7 crossover). The stale
'nufftax needs >=3.12, venv is 3.10' caveat is moot -- the stack now floors
requires-python >=3.12.
Verified by running simulate() for all four instruments (sma, alma,
alma_high_res, hannah); all pass, and alma_high_res completes without OOM.
Closes #128
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JEXzQpvG3QNUdTh6tZcaAE
Four sites claimed TransformerNUFFT is 'pynufft-based' and 'not JIT-friendly'. Both are wrong: the pynufft backend was deleted by PyAutoArray#475, and the nufftax-backed TransformerNUFFT jits fine. The real subtlety, verified 2026-08-23: transform_mapping_matrix takes a PER-CALL xp argument defaulting to np. With xp=jnp it jits and matches the numpy path to 5e-16 relative; with the default xp=np it raises TracerArrayConversionError under jax.jit. That is what the caveat should say. No behaviour change -- these scripts still use TransformerDFT, now with an accurate reason (baseline comparability, not a JIT limitation). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JEXzQpvG3QNUdTh6tZcaAE
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
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.
Fixes the only executable reference to
TransformerNUFFTPyNUFFT, deleted by PyAutoLabs/PyAutoArray#475, left anywhere in the workspace repos.The bug
jax_profiling/dataset_setup/interferometer.py:140named the deleted class. The transformer dict at:137is built eagerly insidesimulate()(:106), before the key lookup — so every instrument raisedAttributeError, not just thealma_high_resconfig that selected the pynufft backend. Reproduced withsimulate('sma'), a DFT dataset. All JAX-profiling dataset setup in this repo was broken.Scripts Changed
jax_profiling/dataset_setup/interferometer.py"nufft_pynufft"dict arm; repointalma_high_resto"nufft"; rewrite the stale rationale commentjax_profiling/jit/interferometer/delaunay.pyjax_profiling/jit/datacube/delaunay.pyNo behaviour change in the
jit/scripts — they still useTransformerDFT.Why
alma_high_resbecomes NUFFT, not DFTBoth objections in the original comment were checked on 2026-08-23:
n_vis · n_pixcrossover measured during the removal work, where the NUFFT is the only feasible path.requires-python = ">=3.12", and nufftax 0.6.1 needs only>=3.11. A 3.10 venv cannot run current autolens at all.The JIT claim, corrected
Four sites called
TransformerNUFFT"pynufft-based" and "not JIT-friendly". Both wrong. The real subtlety, verified here:transform_mapping_matrixtakes a per-callxpargument defaulting tonp. Withxp=jnpit jits and matches the numpy path to 5e-16 relative; with the default it raisesTracerArrayConversionErrorunderjax.jit. That is what the caveat now says.Validation
simulate()run for every instrument key —sma,alma,alma_high_res,hannah— all pass. The eager dict is exactly why a single-instrument check would not have proved the fix.alma_high_rescompletes without OOM on the nufftax path.Notes for the reviewer
release validation FAILED, unrelated to this change) at PR-open; the human explicitly authorized push + PR-open. Merge remains gated.jax_profiling/dataset/no longer reproduce from the current script — regenerating an untouched config (sma) yields different data plus a differingSMALLDATheader stamp. Unfiled; flagged for separate triage.Closes #128