perf(corrfunc): chunk the near-field reduction to lift the memory ceiling

[TobiBu]

Problem

soft_pair_counts_from_topology (the differentiable corrfunc estimator) materialised the full per-pair soft-bin weight tensors at once:

• cross-leaf: w_cross of shape (P, max_leaf, max_leaf, nbins), P = near leaf-pairs
• within-leaf: w_within of shape (L, max_leaf, max_leaf, nbins)

Peak memory scaled as P·max_leaf²·nbins·8 B with no chunking, so it OOM'd above N≈1e4–2e4 (at theta=0, P = all leaf-pairs = O((N/leaf)²)). value_and_grad roughly doubled it.

Fix

Reduce the near field in fixed-size chunks over the leaf-pair axis (both the cross-leaf P axis and the within-leaf L axis) via a rematerialised jax.lax.scan — new _accumulate_block_pairs helper — accumulating into the (nbins,) output so the (chunk, max_leaf, max_leaf, nbins) weight tensor is the only large live intermediate. Memory is now O(chunk_size·max_leaf²·nbins).

Two correctness details:

• jax.checkpoint on the scan body is essential. Without it, reverse-mode would stack every chunk's distance/weight residual over the scan axis and restore the exact O(P) ceiling this removes. With it, backward recomputes each chunk, so peak grad memory ≈ one chunk.
• Single-chunk fast path (q ≤ chunk_size): reproduces the original reduction with no padding/scan, so small-N is bit-identical to the pre-fix code.

chunk_size is a new keyword-only parameter (default 1024) on soft_pair_counts_from_topology and soft_pair_counts — backward-compatible. build_pair_topology and the far-field monopole path are unchanged.

Validation

• Numerical identity (N=1000, θ∈{0,0.3,0.6}, forward + jax.grad): default chunk = 0.0 abs diff vs the pre-fix implementation; scan path (torture-tested at chunk=7/128) matches to ~1e-15 relative (float64 ULP).
• pytest -o addopts="" tests/applications/: 14 passed.
• GPU (A100-40GB, autocvd free GPU):
  • validate_vs_baseline.py --sizes 1000 10000 100000 — no OOM; overlapping-size errors match the committed 40 GB run exactly (estimator counts within 6.08e-16 relative).
  • scaling.py --sizes 10000 100000 1000000 — no OOM: N=1e5 acc 0.51 s / vgrad 0.93 s; N=1e6 acc 7.7 s / vgrad 15.7 s.
• black + basedpyright clean on the changed files.

Bench harness (no estimator-code change)

Two pre-existing, non-estimator OOMs blocked the acceptance commands:

1. validate_vs_baseline.py — the O(N²) brute-force baseline OOMs at 1e5 (that's why the committed validation stopped at 1e4). Added --brute-force-max-n (default 20000); beyond it the θ=0 estimator (near field exact) is used as the reference, so the estimator is still exercised at 1e5. Also enlarged the undersized pair queue (θ=0 at 1e5 makes the whole field near).
2. scaling.py — the traversal preallocated a num_nodes × (1<<15) far-interaction buffer ≈ 8.6 GiB at 1e6, OOMing the build before the estimator ran. Right-sized the caps to 1<<13 (still >40× the observed θ=0.5 peaks); list contents and results unchanged. Also warmed the build_s timing with time_callable and recorded full build stats.

🤖 Generated with Claude Code