[AutoDiff] Autodiff 1: Add baseline adstack regression test for unary_collections

[duburcqa]

Baseline adstack regression test for unary_collections

Test-only PR. Pins reverse-mode AD through a dynamic loop over every unary op currently in NonLinearOps::unary_collections, cross-checked against PyTorch autograd in both f32 and f64.

TL;DR

@qd.kernel
def compute():
    for i in x:
        acc = 0.0
        for j in range(n_iter):
            a = x[i] + qd.cast(j, qd_dtype) * step + offset
            acc += qd_op(a)
        y[None] += acc

a is rebuilt from scratch every iteration of the inner dynamic range-for, so the reverse pass cannot reuse a single BackupSSA-spilled copy — it has to walk the adstack. Any nonlinear unary op dropped from unary_collections causes the multi-iteration variants (n_iter >= 3) to fail with gradients off by large factors; single-iteration variants happen to pass even without promotion, which is noted in the comment.

Why

The existing reverse-mode AD tests in test_ad_*.py exercise the formulas, not the per-iteration spill path. The bug class that motivates this PR is a unary op silently not being in NonLinearOps::unary_collections (the set that drives AdStackAllocaJudger's decision to promote the operand alloca): without promotion, BackupSSA spills the operand to a single plain alloca that every forward iteration overwrites, so every backward step reads the last-iteration value. Gradients come out silently wrong at n_iter >= 3.

Pinning this explicitly with a per-op parametrize means (a) every op currently in unary_collections has a failing-case regression, and (b) subsequent PRs in the chain get a ready-made slot to drop new ops into.

Surface

Single new test file: tests/python/test_adstack.py.

• Module-level _UNARY_OPS_PARAMS tuple: (op_name, step, offset) per op, grouped by domain (real / positive-subunit) with comments explaining the (step, offset) choices.
• _run_unary_loop_carried(qd_dtype, op_name, step, offset, x_val, n_iter, rel_tol) helper: builds the kernel, runs forward + backward, compares against a PyTorch reference in the matching dtype.
• test_adstack_unary_loop_carried — f32, rel=1e-4, decorated with require=qd.extension.adstack.
• test_adstack_unary_loop_carried_f64 — f64, rel=1e-12, decorated with require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64.

Both iterate the same _UNARY_OPS_PARAMS × x_val ∈ {0.001, 0.15, 0.26, 0.399} × n_iter ∈ {1, 3, 10} matrix. 0.26 (rather than 0.25) avoids x[3] + offset = 0 at j=0, where abs's derivative is undefined. abs/sin/cos use a (step, offset) that makes the operand cross zero across iterations so their gradient signs actually vary (otherwise they'd pass even if dropped from the set).

Why two separate tests instead of parametrizing dtype

qd.cast(py_float_literal, qd.f64) bakes the literal at default_fp first, then casts to f64 — so qd.cast(0.3, qd.f64) is not bit-equal to IEEE-754 0.3. The only way to get truly tight (~1e-12) f64 tolerance is to decorate the test with default_fp=qd.f64 so the frontend treats Python floats as f64 from the start. That forces the two tests to be separate at the decorator level; the shared body lives in _run_unary_loop_carried.

Tests

All parametrize entries pass across arm64 / Metal / Vulkan (local) and x86_64 / AMDGPU / Vulkan (amdgpu).

Stack

Autodiff 1 of 13. Followed by #501 (tan).

[hughperkins]

@claude review

[hughperkins]

@codex review

[hughperkins]

Opus summary:

Summary

Adds a baseline regression test for reverse-mode AD (adstack) through dynamic loops
whose loop-variant operand feeding a non-linear unary op is loaded but never stored back
in the inner body. This pattern exercises the AdStackAllocaJudger promotion path that
depends on NonLinearOps::unary_collections membership — without promotion, BackupSSA
spills the operand into a single plain alloca that every forward iteration overwrites,
producing silently wrong gradients on the backward pass.
The test parametrizes over the unary ops currently in unary_collections
(abs, sin, cos, asin, acos, log, sqrt) and cross-validates the
Quadrants gradient against PyTorch autograd.

Strengths

• Targets a specific, subtle bug vector. The kernel structure (redefine a from
  scratch each iteration, no LocalStore(a) after LocalLoad(a)) isolates
  unary_collections as the sole promotion trigger, so dropping any op from that set
  causes a test failure.
• Well-documented. The long comment explains why the test is shaped the way it is,
  what breaks if the code under test regresses, and how follow-up PRs will extend it.
• Cross-validates against PyTorch autograd, which is a reliable oracle for gradient
  correctness.
• Designed for incremental extension. The comment explicitly calls out future
  parametrize entries for tan, tanh, exp, and rsqrt in follow-up PRs.
• Proper guards. Uses @pytest.mark.needs_torch, require=qd.extension.adstack,
  and ad_stack_experimental_enabled=True so the test only runs in capable environments.

Weaknesses / things to consider

• Single input value. Only tests x_val = 0.5. Some of these ops have interesting
  behaviour near domain boundaries (e.g. acos/asin near ±1, log/sqrt near 0).
  A second input value closer to a boundary would increase confidence.
• Single field size. The dense axis is size 1. A slightly larger field (e.g. 4)
  would verify that per-element gradient accumulation works correctly across multiple
  elements.
• Fixed loop count. Only range(3) is tested. Adding a second count (e.g. 1 or 10)
  would guard against off-by-one issues in the reversed loop.
• No forward-pass assertion. The test checks gradients but not that the forward
  output y[None] matches the PyTorch reference. A one-line assert would be cheap
  insurance.
• Loose tolerance. rel=1e-3 is generous for float32 arithmetic that involves at
  most 3 loop iterations. rel=1e-5 would likely still pass and catch more drift.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 564dced5c6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[hughperkins]

Please could you address Opus's points:

• Single input value. Only tests x_val = 0.5. Some of these ops have interesting
  behaviour near domain boundaries (e.g. acos/asin near ±1, log/sqrt near 0).
• A second input value closer to a boundary would increase confidence.
  Single field size. The dense axis is size 1. A slightly larger field (e.g. 4)
  would verify that per-element gradient accumulation works correctly across multiple
  elements.
• Fixed loop count. Only range(3) is tested. Adding a second count (e.g. 1 or 10)
  would guard against off-by-one issues in the reversed loop.
• No forward-pass assertion. The test checks gradients but not that the forward
  output y[None] matches the PyTorch reference. A one-line assert would be cheap
  insurance.
• Loose tolerance. rel=1e-3 is generous for float32 arithmetic that involves at
  most 3 loop iterations. rel=1e-5 would likely still pass and catch more drift.

[hughperkins]

this doc is not very readable. might be good to express it in python terms if possible.

[hughperkins]

proposal from opous:

      # Regression test for reverse-mode AD through a dynamic loop that applies a non-linear unary op
      # to a value recomputed each iteration from a field read and the loop index. The key detail is
      # that `a` is rebuilt from scratch every iteration (`x[i] + j * 0.1`) rather than accumulated,
      # so the AD system must save a per-iteration copy of `a` for the backward pass. Without that,
      # every backward step replays the *last* forward value of `a`, producing silently wrong
      # gradients.
      #
      # Each parametrized op (abs, sin, cos, …) must be individually recognised as non-linear for this
      # saving to kick in; dropping any one of them causes the test to fail. Follow-up PRs extend the
      # parametrize list as they add new ops (tan) or fix ops that had buggy reverse formulas
      # (tanh/exp) or were missing entirely (rsqrt).

[duburcqa]

Addressed all five points by hardening the test:

• Multiple x_val ([0.2, 0.3]): one small sample for log/sqrt/rsqrt sensitivity near zero and one closer to the asin/acos domain boundary. The 0.05 step keeps x_val + (n_iter-1)*0.05 + (n-1)*0.05 below 1.0 for every op's domain.
• Multiple n_iter ([1, 3, 10]): 1 smoke-tests the single-push adstack path, 10 stresses accumulation over many iterations.
• Field size 4 (was 1): catches per-element gradient accumulation regressions. Each element takes a different x_val offset, so every test case cross-validates gradients at multiple input points at once.
• Forward-pass assertion added: assert y[None] == test_utils.approx(y_t.item(), rel=1e-5) alongside the per-element grad check.
• Tolerance tightened from 1e-3 to 1e-5 on both the forward and backward assertions.

The parametrize list is unchanged (still the 7 ops safe on main); follow-up PRs #501/#502/#503 extend it with tan, tanh/exp, and rsqrt.

[hughperkins]

why split this out? why not fuse it with the earlier test, and simply paramterize over some negative values too? (sin, cos both work with think? For values that arent compatible, you can check that both epxected and actual are nan I guess?)

[duburcqa]

It is explained in the description of the unit test. Unless you disagree with it.

[hughperkins]

so you are iterating over j, and using j in the formula for a? why cant we do the same thing in the first unit test, and unify the two tests?

[hughperkins]

add a negative value, and add abs to the list of ops?

[hughperkins]

Reran opus:

Weaknesses / things to consider

• No domain-boundary inputs. Input values 0.2/0.3 stay safely in the interior
  of every op's domain. Values near boundaries (e.g. acos/asin near ±1, log/sqrt
  near 0) would test numerical stability of the gradient formulas.
• Docstring still references internal compiler concepts (NonLinearOps::unary_collections,
  "promotion", "plain scratch slot"). The opening sentence is now plain English, but the
  detail could be further simplified for someone reading only the Python test.
• abs test uses hand-computed reference instead of PyTorch. This is arguably fine
  (the formula is trivial), but using torch.abs for consistency would make it obvious
  that the same oracle strategy applies everywhere.

[hughperkins]

Please address Opus's last comments, ie:

• No domain-boundary inputs. Input values 0.2/0.3 stay safely in the interior
  of every op's domain. Values near boundaries (e.g. acos/asin near ±1, log/sqrt
  near 0) would test numerical stability of the gradient formulas.
• Docstring still references internal compiler concepts (NonLinearOps::unary_collections,
  "promotion", "plain scratch slot"). The opening sentence is now plain English, but the
  detail could be further simplified for someone reading only the Python test.
• abs test uses hand-computed reference instead of PyTorch. This is arguably fine
  (the formula is trivial), but using torch.abs for consistency would make it obvious
  that the same oracle strategy applies everywhere.

[claude]

LGTM — all prior concerns addressed: sign-crossing operand, redundant kwarg removed, and docstring invariant scoped to n_iter >= 3.

Extended reasoning...

Overview

This PR adds a single new test file () with a parametrized regression test for reverse-mode AD through dynamic loops whose loop-variant operand feeds a non-linear unary op. The test cross-validates Quadrants gradients against PyTorch autograd across 7 ops, 4 x_val boundary samples, and 3 n_iter counts.

Security risks

None. This is a test-only addition with no production code changes.

Level of scrutiny

Low — test-only change, no logic or API surface affected. The mathematical correctness of the test design is the main concern, and three rounds of review have thoroughly vetted it.

Other factors

All prior inline review concerns have been resolved: the case now uses to cross zero (making correct and wrong gradients distinguishable), the redundant kwarg was dropped, the misleading comment formula was removed, and the docstring now correctly scopes the regression invariant to while explaining that is a single-push smoke test. The -without- pattern is confirmed as the repo convention. No outstanding comments remain unaddressed.

[hughperkins]

Updated opus description:

Summary

Adds a single new regression test, tests/python/test_adstack.py, that pins reverse-mode AD behavior for loop-carried unary ops.
The test cross-checks d/dx Σ_j op(x + j*step + offset) against PyTorch autograd for a parametrized set of unary ops (sin, cos, abs, log, sqrt, asin, acos),
across n_iter ∈ {1, 3, 10} and x_val ∈ {0.001, 0.15, 0.26, 0.399}. The intent is to catch regressions where a unary op is dropped from the AD transform's per-iteration
operand set: the operand then falls back to a single overwritten slot, and the reversed loop reads the last-iteration value for every backward step, producing wrong gradients
at n_iter ≥ 3.

What's in the PR

Net diff vs origin/main: 1 file, +89 / -0 — purely additive.

• New file: tests/python/test_adstack.py (+89 lines).
• Per-op (step, offset) chosen so the operand stays inside the op's domain for every (x_val, n_iter) combination, while still crossing zero for abs/sin/cos so the
  per-iteration sign actually varies (otherwise abs's piecewise-constant derivative would make the test pass trivially).
• x_val samples cover both interior points and domain edges (0.001 near 0 for log/sqrt; 0.399 near 1 for asin/acos).
• Tolerance: rel=1e-4 (loosened from a tighter initial value to accommodate f32).
• Gated by @pytest.mark.needs_torch and @test_utils.test(require=qd.extension.adstack).
  Commits (4 + 1 merge):
• 848ae6bfb Add baseline adstack regression test for unary_collections.
• 15befa000 Simplify adstack unary test docstring per review.
• 8e99b4631 Rephrase adstack unary test docstring to drop remaining compiler-internal terms.
• 13787f5b6 Rephrase adstack unary test docstring as user-facing and loosen tolerance to rel=1e-4 for f32.
• 2d016ffb4 Merge main into branch.

Good points

• Targeted regression coverage. Pins a specific, known failure mode (loop-variant unary op missing from the AD per-iteration operand set) that wouldn't be caught by
  single-iteration tests.
• Cross-checked against PyTorch rather than against a hand-derived analytical gradient, so the oracle is independently maintained.
• Domain- and edge-aware parametrization. x_val deliberately probes near-singular points of each op's gradient (log/sqrt near 0, asin/acos near 1), and 0.26 is
  chosen instead of 0.25 specifically to avoid abs's undefined derivative at 0.
• Sign-crossing operand for abs/sin/cos ensures the per-iteration sign varies, so the test isn't trivially satisfied by a constant-derivative shortcut.
• n_iter = 1 included as a sanity baseline alongside the n_iter ∈ {3, 10} cases that actually exercise the loop-carried path.
• Self-documenting: the docstring explains both the user-facing intent and the internal AD mechanism the test pins, plus how to extend the parametrize list when a new
  unary op becomes differentiable through dynamic loops.
• Clean, additive diff. Now that main is merged in, the PR is exactly one new file — easy to review, zero risk of touching unrelated code.

Bad points / concerns

• Test is a pin, not a fix. It locks in current behavior but doesn't itself address any bug; if the underlying AD operand-set is fragile, this test surfaces regressions
  but doesn't prevent the class of issue.
• Single dtype (f32) only. No f64 variant, despite the tolerance discussion implying f32 was chosen specifically. An f64 parametrization with a tighter tolerance would
  strengthen the oracle.
• Tolerance was loosened (rel=1e-4). Reasonable for f32, but worth confirming the loosening was driven by genuine f32 round-off rather than masking a real numerical
  issue.
• Parametrize matrix is 3 × 4 × 7 = 84 cases, all PyTorch-gated. Not large in absolute terms, but every new unary op added per the docstring's extension instructions
  multiplies by 12, so the suite will grow if followed literally.
• Hard-coded n = 4 and stride 0.05 for both the Quadrants field initialization and the PyTorch reference; named constants or a small helper would make the coupling
  between the two sides more obvious and harder to desync.
• Merge commit rather than rebase. Minor style point — the branch was brought up to date via git merge main rather than rebased, so the final history includes a merge
  commit. Depending on repo conventions, a rebase + force-push might be preferred.

[hughperkins]

can you comment on:

Single dtype (f32) only. No f64 variant, despite the tolerance discussion implying f32 was chosen specifically. An f64 parametrization with a tighter tolerance would
strengthen the oracle.

please

[claude]

LGTM — all raised issues have been addressed, including the f64 tolerance floor fix.

Extended reasoning...

Overview

Single new test file tests/python/test_adstack.py (+107 lines, no production code changes). It pins reverse-mode AD correctness for loop-carried unary ops by cross-checking gradients against PyTorch autograd across seven ops, four x_val samples, and three n_iter counts in both f32 and f64.

Security risks

None. This is a test-only addition with no new executable production paths.

Level of scrutiny

Low. The change is purely additive, confined to the test suite, and has no impact on any production or library code. The correctness of the test itself is the only concern.

Other factors

The PR went through several rounds of review feedback, all of which have been addressed in the current revision:

• The abs parametrize case is no longer degenerate: ("abs", 0.3, -0.4) ensures the operand crosses zero so per-iteration signs vary.
• The f64 tolerance concern (my 2026-04-21 inline comment) is fixed: assertions now use pytest.approx directly, bypassing the test_utils.approx 1e-6 floor, with an explanatory inline comment.
• The redundant ad_stack_experimental_enabled=True kwarg has been removed.
• The docstring accurately scopes the multi-iteration regression invariant to n_iter >= 3 and notes n_iter = 1 is a single-push smoke test.
• No bugs were found by the automated bug hunting system.

[hughperkins]

Checklist:

• just a bug fix => doesn't change existing API or usage
  • no need for doc changes
• doesn't add significantly more code around diverse files

=> ok to merge