Factor shared scaffolding out of the allocating Enzyme rules

[CarloLucibello]

Follow-up refactor after #735 (merged). No behavior change — purely reduces boilerplate in the Enzyme extension.

Motivation

The 8 EnzymeRules for allocating ops (y = f(...)) — imrotate, upsample_nearest, upsample_linear, grid_sample, ctc_loss, batchnorm, unfold, fold — each repeated the same augmented/reverse scaffolding: run the forward, allocate an output shadow, then on reverse loop over the batch width feeding the shadow to the op's ∇ function and accumulating into each differentiable input. Adding an op meant copy-pasting that plumbing (and it was exactly that copy-paste that hid the width-1 _enz_slice scalar-indexing bug fixed in #735).

Change

Four small helpers capture the shared structure:

• _enz_alloc_augmented(config, y, extra) — the standard augmented pass (primal/shadow per config, cache (shadow, extra...)).
• _enz_alloc_reverse!(config, shadow, args, backward) — the reverse loop; backward(dy) returns one cotangent per arg.
• _enz_accum! — a Base.tail recursive unroll of the per-input accumulation, keeping it type-stable (a zip over the heterogeneous arg/grad tuples would dynamically dispatch).
• _enz_active_reverse! — the scalar Active-return variant (ctc_loss).

Each rule now shrinks to its forward call plus a one-line backward closure, so only the per-op specifics (which ∇, which inputs, what to cache) remain visible. The now-unused _enz_pairs helper is removed. Signatures stay per-rule (needed for dispatch, kwargs defaults, and the Const{<:DenseConvDims} constraints). Scope is deliberately limited to these allocating rules; the in-place conv!/pool!/softmax!/gather!/scatter!/_dropout! rules and batched_mul are left untouched (their overwritten-based caching is subtle and varies per rule).

Net −7 lines despite adding the helper infrastructure, and each of the 8 rules is roughly halved.

Verification (behavior-preserving)

• CPU testsuites unchanged: rotation 347/347, upsample 74/74, fold 67/67, sampling 30/30, ctc 5/5.
• EnzymeTestUtils.test_reverse across all activity combinations including BatchDuplicated (width > 1): imrotate, upsample_nearest, upsample_linear, unfold (123/123), fold (147/147) all pass. grid_sample reproduces the same pre-existing finite-difference discrepancies (grid gradient under 0-padding) byte-for-byte against the pre-refactor code (78 pass / 6 fail on both).
• GPU batchnorm (g/b/x) and grid_sample (input/grid) still match Zygote.
• Untouched in-place rules pass: cpu/pooling 1402/1402, cpu/conv 4117 pass.

🤖 Generated with Claude Code