refactor: chunk the kernel-CDF forward transform (remove the O(M×N) memory wall) #376

Overview

The kernel-CDF forward transform (rectangular_kernel.py, #373/PR#374) broadcasts queries × points into an (M, N, 2) array — ~60 GB at production imaging scale (M≈246k, N≈15.4k; observed OOM 2026-07-10). Chunk the query dimension so peak memory is O(block×N) with numerics identical to float precision. Refactor invariant: all public behaviour unchanged; witnesses: test_autoarray suite + both jax_grad certification scripts unchanged.

Plan

• Chunk F_raw over queries in create_transforms_kernel: jax path via lax.map over fixed-size blocks (pad M to a block multiple, slice back; AD flows through blocks); numpy path via a block loop. Block size an internal constant (~1024–4096, tuned for ~O(100 MB) peak at production scale).
• Unit test: chunked forward == unchunked reference on random inputs, incl. non-multiple-of-block M (numpy path).
• Re-run jax_grad/imaging_pixelization.py + interferometer.py — all variants pass with unchanged values (strict FD, eager-vs-JIT, parity).
• Demonstrate the previously-OOM scale runs within laptop memory (direct probe at M≈246k, N≈15.4k).
• Update the O(M×N) caveat in autolens_workspace_developer/jax_profiling/gradient/README.md (workspace follow-up commit).

Branch Survey

PyAutoArray on main (post-#374/#375 merge), clean. Branch: feature/kernel-forward-chunking. Worktree: ~/Code/PyAutoLabs-wt/kernel-forward-chunking/.
Note: active.md still carries rect-adapt's vestigial PyAutoArray claim (its PR #375 merged today; entry awaiting retirement by its owning session) — proceeding under the same human coordination that directed this follow-up.

Original Prompt

PyAutoMind/refactor/autoarray/kernel_forward_chunking.md (verbatim in repo).

🤖 Generated with Claude Code