Skip to content

fix: drop deleted TransformerNUFFTPyNUFFT from jax_profiling dataset setup - #129

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/pynufft-removal-residue-phase-1
Aug 23, 2026
Merged

fix: drop deleted TransformerNUFFTPyNUFFT from jax_profiling dataset setup#129
Jammy2211 merged 2 commits into
mainfrom
feature/pynufft-removal-residue-phase-1

Conversation

@Jammy2211

Copy link
Copy Markdown
Contributor

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:140 named the deleted class. The transformer dict at :137 is built eagerly inside simulate() (:106), before the key lookup — so every instrument raised AttributeError, not just the alma_high_res config that selected the pynufft backend. Reproduced with simulate('sma'), a DFT dataset. All JAX-profiling dataset setup in this repo was broken.

Scripts Changed

Script Change
jax_profiling/dataset_setup/interferometer.py Drop the "nufft_pynufft" dict arm; repoint alma_high_res to "nufft"; rewrite the stale rationale comment
jax_profiling/jit/interferometer/delaunay.py Correct three stale pynufft/JIT claims (prose only)
jax_profiling/jit/datacube/delaunay.py Correct one stale pynufft/JIT claim (prose only)

No behaviour change in the jit/ scripts — they still use TransformerDFT.

Why alma_high_res becomes NUFFT, not DFT

Both objections in the original comment were checked on 2026-08-23:

  • "DFT would need a dense (n_vis × n_real_space) matrix = ~20GB and OOM on a 15GB laptop"still true, and it rules DFT out. 5000 vis × 512×512 = 1.31e9, far above the ~1e7 n_vis · n_pix crossover measured during the removal work, where the NUFFT is the only feasible path.
  • "nufftax requires Python >= 3.12 (PyAutoGPU venv is 3.10)"obsolete. PyAutoArray, PyAutoLens and PyAutoNerves all declare 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_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 it raises TracerArrayConversionError under jax.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_res completes without OOM on the nufftax path.
  • Alias-aware AST sweep: 0 executable references to the deleted class remain in this repo.
  • Generated datasets were reverted — this PR touches source only.

Notes for the reviewer

  • Heart was RED (release validation FAILED, unrelated to this change) at PR-open; the human explicitly authorized push + PR-open. Merge remains gated.
  • Pre-existing and not addressed here: the committed datasets under jax_profiling/dataset/ no longer reproduce from the current script — regenerating an untouched config (sma) yields different data plus a differing SMALLDAT header stamp. Unfiled; flagged for separate triage.

Closes #128

Jammy2211 and others added 2 commits August 23, 2026 17:43
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release Pending release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: TransformerNUFFTPyNUFFT AttributeError breaks all dataset setup

1 participant