Skip to content

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

Description

@Jammy2211

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions