fix: TransformerNUFFTPyNUFFT AttributeError breaks all dataset setup

[Jammy2211]

Overview

TransformerNUFFTPyNUFFT was deleted from PyAutoArray by PyAutoLabs/PyAutoArray#475 (2026-08-22), but jax_profiling/dataset_setup/interferometer.py still references it. Because the transformer lookup dict is built eagerly inside simulate(), every instrument raises AttributeError — not just the one config that selects the pynufft backend. All JAX-profiling dataset setup in this repo is currently broken.

This is Phase 1 of 3 cleaning up residue from the pynufft removal, whose workspace tier was scoped to autolens_workspace / autolens_workspace_test only and never swept the sibling repos. It is the only executable reference to the deleted class anywhere; the other two phases are documentation and CI.

Plan

• Reproduce the failure on clean main (confirmed 2026-08-23 via simulate('sma') — a DFT dataset that still fails, proving the break is not confined to the pynufft config).
• Remove the "nufft_pynufft" arm from the transformer dict.
• Repoint the alma_high_res config to the nufftax-backed TransformerNUFFT, rather than DFT — its ~20GB dense-matrix OOM constraint is still real.
• Rewrite the stale rationale comment to record why the config is NUFFT rather than deleting the reasoning.
• Verify by running simulate() across every instrument key, not just the changed one.

Detailed implementation plan

Affected Repositories

• autolens_workspace_developer (primary)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_developer	main	clean

No worktree claims; worktree_check_conflict exit 0.

Suggested branch: feature/pynufft-removal-residue

The failure

jax_profiling/dataset_setup/interferometer.py:140
    "nufft_pynufft": al.TransformerNUFFTPyNUFFT,

AttributeError: module 'autolens' has no attribute 'TransformerNUFFTPyNUFFT'

The dict at :137 is constructed inside simulate() (:106) before the key lookup at :141, so the attribute is evaluated on every call regardless of which instrument is requested. Reproduced with simulate('sma'), whose config sets transformer_class: "dft".

Implementation Steps

1. jax_profiling/dataset_setup/interferometer.py:137-141 — drop the "nufft_pynufft": al.TransformerNUFFTPyNUFFT, entry, leaving "dft" and "nufft".
2. jax_profiling/dataset_setup/interferometer.py:76 — change alma_high_res's "transformer_class" from "nufft_pynufft" to "nufft".
3. jax_profiling/dataset_setup/interferometer.py:65-69 — rewrite the rationale comment. Both of its claims were checked on 2026-08-23:
   • "DFT would need a dense (n_vis x n_real_space) matrix = ~20GB and OOM on a 15GB laptop" — still true, and it is why this config must not become DFT. 5000 vis x 512x512 = 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.

Key Files

• jax_profiling/dataset_setup/interferometer.py — simulate() at :106; instrument configs from :60; transformer dict at :137.

Verification

• Call simulate() for every instrument key and require no AttributeError. The eager dict is precisely why testing one instrument would not prove the fix.
• Run the alma_high_res path specifically and confirm it produces a dataset rather than OOMing.
• Re-run an alias-aware attribute sweep over the post-fix tree; require zero executable references to TransformerNUFFTPyNUFFT. Do not treat the inventory in this issue as proof of completeness.

Why this rotted unnoticed

This repo has no smoke coverage — the same root cause recorded in the repo's broader stale_api_rot_audit prompt (56 stale symbols, scanned 2026-08-04). TransformerNUFFTPyNUFFT only became stale on 2026-08-22, so it is absent from that inventory: same repo and same class of rot, different instance. If a minimal smoke tier is added there, this file is a strong candidate for it.

Original Prompt

Click to expand starting prompt

Phase 1: fix the pynufft AttributeError breaking all jax_profiling dataset setup

Type: maintenance
Target: autolens_workspace_developer
Repos:

• @autolens_workspace_developer
  Difficulty: low
  Autonomy: supervised
  Priority: normal
  Status: draft
  Filed: 2026-08-23

Phase 1 of 3. Parent: pynufft_removal_downstream_residue.md (full evidence,
provenance and out-of-scope list live there). Phases are independent — no
ordering constraint, no library-first gate, because no PyAuto* library source
changes are involved.

The break (reproduced 2026-08-23)

jax_profiling/dataset_setup/interferometer.py

:140        "nufft_pynufft": al.TransformerNUFFTPyNUFFT,

AttributeError: module 'autolens' has no attribute 'TransformerNUFFTPyNUFFT'

TransformerNUFFTPyNUFFT was deleted by @PyAutoArray#475 (2026-08-22). The
dict at :137 is built eagerly inside simulate() (:106), before the key
lookup — so every instrument raises, not only the alma_high_res config at
:76 that selects "nufft_pynufft". Confirmed by calling simulate('sma'), a
DFT dataset, which still fails. All JAX-profiling dataset setup in this repo
is currently broken.

This is the only executable reference to the deleted class in any repo.

Task

Drop the "nufft_pynufft" arm of the dict and repoint the alma_high_res
config's transformer_class to "nufft" (al.TransformerNUFFT, nufftax-backed).

This is a decision, not a mechanical rename. Both objections in the :65-69
comment were checked (2026-08-23):

• "DFT would need a dense (n_vis x n_real_space) matrix = ~20GB and OOM on a
  15GB laptop" — still true, and it rules DFT out here. 5000 vis x 512x512
  = 1.31e9, far above the ~1e7 n_vis * n_pix crossover measured in
  remove_pynufft_legacy_transformer.md, 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.

Rewrite that comment to record the new rationale rather than deleting it — the
OOM constraint is still the reason this config is not DFT.

Why it went unnoticed

This repo has no smoke coverage, the same root cause recorded in
draft/maintenance/autolens_workspace_developer/stale_api_rot_audit.md
(Status: formalised). That prompt's alias-aware scan ran 2026-08-04 and found 56
stale symbols; TransformerNUFFTPyNUFFT only became stale on 2026-08-22, so it
is absent from that inventory — same repo, same class of rot, different
instance. If a minimal smoke tier is added under that prompt,
dataset_setup/interferometer.py is a strong candidate for it.

Acceptance

• simulate() runs for every instrument key, not just the changed one — the
  eager dict is exactly why a single-instrument check would miss a regression here.
• The alma_high_res path produces a dataset rather than OOMing.
• No executable reference to TransformerNUFFTPyNUFFT remains in the repo
  (re-run an alias-aware attribute sweep; do not trust the inventory above as proof).
• The alma_high_res rationale comment reflects the decision actually taken.

[Jammy2211]

Branch correction: the task worktree helper derives the branch from the task name, so the working branch is feature/pynufft-removal-residue-phase-1, not feature/pynufft-removal-residue as the issue body states.

Worktree: ~/Code/PyAutoLabs-wt/pynufft-removal-residue-phase-1
Registered in PyAutoMind/active.md as pynufft-removal-residue-phase-1 (status: workspace-dev).

[Jammy2211]

Shipped

Merged in #129 (2026-08-23).

The fix. Dropped the "nufft_pynufft" arm naming the deleted al.TransformerNUFFTPyNUFFT, and repointed alma_high_res to the nufftax-backed TransformerNUFFT. The dict at :137 is built eagerly inside simulate(), so this had been breaking every instrument, not just the one config that selected pynufft.

Why NUFFT and not DFT. The comment's ~20GB dense-matrix OOM objection is still real (5000 vis x 512x512 = 1.31e9, far above the ~1e7 crossover). Its second objection — "nufftax requires Python >= 3.12 (PyAutoGPU venv is 3.10)" — is obsolete: the whole stack declares requires-python = ">=3.12" and nufftax 0.6.1 needs only >=3.11.

Bonus finding. Four sites in jax_profiling/jit/ claimed TransformerNUFFT is "pynufft-based" and "not JIT-friendly". Both wrong. Verified: transform_mapping_matrix takes a per-call xp defaulting to np; with xp=jnp it jits and matches numpy to 5e-16, with the default it raises TracerArrayConversionError. Corrected, no behaviour change.

Validation. simulate() run for all four instruments (sma, alma, alma_high_res, hannah) — all pass, alma_high_res without OOM. Alias-aware AST sweep: 0 executable references to the deleted class remain anywhere.

Merged over two known reds, both human-authorized: Heart RED (release validation FAILED, unrelated) and — for the sibling assistant PRs — a pre-existing wiki-currency red that exists on main independently of this work.

Two things found and deliberately left unfiled

1. This repo has no test CI (one Copilot workflow). That is why this rotted unnoticed, and why fix: drop deleted TransformerNUFFTPyNUFFT from jax_profiling dataset setup #129 merged with no checks. The broader stale_api_rot_audit prompt records the same root cause.
2. The committed datasets under jax_profiling/dataset/ no longer reproduce from their own scripts — regenerating an untouched config (sma) yields different data plus a differing SMALLDAT header stamp. Pre-existing, unrelated to pynufft.