[GeophyAI]

Fixes the cross-shot coupling of ViscoAcoustic with per-shot batched models (B, nz, nx) found while answering #123.

Problem

prepare_models froze the fractional exponent at gbar = mean(gamma) over the whole batch, so each shot's |k|^(2*gbar+1) / |k|^(2*gbar+2) filters depended on the other shots' Q. With Q different across the batch, batched vs one-shot-at-a-time differed by ~1.7e-3 (data) and ~8e-4 (gradient), on both impl='eager' and impl='c'.

Fix

• ViscoAcoustic._average_exponent: one gbar per batch element for per-shot models ((B, 1, 1, 1)), each the same reduction a single-shot run takes. Shared models keep the scalar; a shared model that the C path repeats across the batch (all per-element means equal) keeps the historical whole-batch mean bit-for-bit.
• c_eq_aux: builds (B, nz, nx) D_loss / D_frac grids when the exponents differ (D_k2 stays (nz, nx)).
• CUDA (visco_acoustic2d/kernels.cuh, backward.cu): the grid check accepts (nz, nx) or (B, nz, nx); lop_into / cmul_grid index the grid modulo its own size (unchanged nz*nx for a shared grid), with a check that a grid is either one plane or the whole batch.

Tests

New test/test_visco_acoustic_batched_models.py (Q different per shot):

path	batched vs loop: data	vp / Q gradient (rel)
before	1.7e-3	8e-4
c, full	bit-identical	bit-identical
c, ckpt chunk / recursive	bit-identical	2e-7 – 3e-6
eager	bit-identical	2e-7 – 6e-7

Plus a CPU unit test of the per-element exponent and a check that a shared model on impl='c' keeps (nz, nx) grids and equals an identical per-shot batch bit-for-bit.

Shared-model A/B against dev (c full / c default / eager / eager compiled × the three phase_shift / amplitude_damping combinations, V100, core rebuilt from this branch): all records bit-identical; the gradients that differ between dev and this branch are exactly those that already differ between two dev runs (atomic accumulation), at the same magnitude.

Existing test_visco_acoustic.py, test_visco_acoustic_cuda.py, test_visco_acoustic_jax.py, test_elastic_batched_models.py, test_backward_scratch_from_pool.py: 107 passed.

🤖 Generated with Claude Code