[AutoDiff] Autodiff 6: Adstack regression tests

[duburcqa]

Adstack regression tests

Test-only PR. Pins baseline reverse-mode AD behaviour through dynamic and static for-loops with the adstack extension. Complements Autodiff 1's unary-op test with coverage of the linear-kernel path, which has a different failure signature and different structural invariants.

TL;DR

Three structural tests, each with an f32 and an f64 variant:

1. test_adstack_basic_gradient(_f64) — loop-carried state mutation. v = v * 0.95 + 0.01 iterated n_iter ∈ {1, 3, 10} times. Gradient dy/dx[i] = 0.95**n_iter, independent of x[i], so a broken push/load/pop pattern would still happen to produce the right answer (the backward chain adj(v_prev) = 0.95 * adj(v_next) only uses the compile-time constant 0.95). What this test pins is structural: the backward compiler must successfully reverse a dynamic range(n_iter).

2. test_adstack_basic_gradient_negative — the negative counterpart. Decorated with @test_utils.test(ad_stack_experimental_enabled=False); asserts compute.grad() raises qd.QuadrantsCompilationError("Cannot use non static range in Backwards mode") deterministically. Pins the exact compile-time rejection.

3. test_adstack_sum_linear(_f64) — 2×2×3 parametrize matrix: (static vs dynamic loop) × (constant vs loop-varying coefficient) × (n_iter ∈ {1, 3, 10}). Linear accumulation y = sum_j v * coeff_j; v is loop-invariant. Proves that enabling the adstack extension does not silently regress linear reverse-mode AD in any of the four shape combinations.

Why

Autodiff 1 pins reverse-mode AD for nonlinear unary ops in dynamic loops — the shape where the per-iteration operand has to be replayed from the adstack. The class of bugs it catches: a unary op missing from unary_collections → operand falls back to a single spill slot → last-iteration value read for every backward step.

This PR pins the orthogonal shape: linear loop bodies. The failure modes here are different:

• The adstack is structurally required to reverse a dynamic range at all — without it, the backward compiler flat-out refuses to build the kernel. That's what test_adstack_basic_gradient_negative pins.
• Once the adstack enables the reversal, value-correctness of the spilled state depends on push/load/pop timing relative to the loop control flow. The linear-body tests do not exercise per-iteration value correctness of the stored v (the coefficients are compile-time constants), but they catch structural bugs that would stop the backward kernel from running at all — mis-sized slab, wrong iteration direction, off-by-one in the pop count on the reversed loop, and similar.

Both classes of test are needed: a unary-ops regression would slip past the linear tests, and a structural regression (e.g. an off-by-one that nonetheless happens to produce the right gradient for non-linear unary ops through compensation) would slip past the unary tests.

The f64 design

Each test is a thin wrapper around a helper that takes qd_dtype and rel_tol:

def _run_basic_gradient(qd_dtype, n_iter, rel_tol, shall_not_pass): ...
def _run_sum_linear(qd_dtype, use_static_loop, use_varying_coeff, n_iter, rel_tol): ...

The f32 and f64 tests decorate with different @test_utils.test(...) modes:

• f32: @test_utils.test(require=qd.extension.adstack), rel=1e-6 (tight for a strictly linear kernel; catches a new fp op introduced in the backward).
• f64: @test_utils.test(require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64), rel=1e-14.

default_fp=qd.f64 at decorator time is load-bearing: Python float literals in the kernel body (0.95, 0.01, etc.) get baked at whatever default_fp says, so 0.95 becomes an f32-quantized constant if the decorator uses the default default_fp=qd.f32, and the resulting gradient is not actually f64-precision. Splitting into two tests with different decorators is the only way to get truly tight f64 tolerance; the shared body lives in the helper.

Coverage matrix

Test	Dynamic loop	Static loop	Loop-carried state	Check
test_adstack_basic_gradient	✓		✓	value match against 0.95**n_iter
test_adstack_basic_gradient_negative	✓		✓	compile-time rejection
test_adstack_sum_linear[static, const]		✓		value match against n_iter
test_adstack_sum_linear[static, varying]		✓		value match against sum(a+1 for a)
test_adstack_sum_linear[dyn, const]	✓			same
test_adstack_sum_linear[dyn, varying]	✓			same

Stack

Autodiff 6 of 13. Based on #496 (alloca placement). Followed by #534 (header size).

[chatgpt-codex-connector]

💡 Codex Review

https://github.com/Genesis-Embodied-AI/quadrants/blob/c7d23e5698683e4b69b4cdfefa9ce85d6c94bb42/transforms/auto_diff.cpp#L1261-L1266
Add tan to nonlinear-op stack detection

Now that UnaryOpType::tan is differentiated, operands feeding tan in loop-carried locals must be treated as nonlinear for ad-stack backup, but NonLinearOps::unary_collections (used by AdStackAllocaJudger::visit(UnaryOpStmt)) still omits tan. In reverse mode with adstack, this means some allocas used by tan are left as plain locals, so overwritten primal values are not preserved across iterations and gradients can be computed from the wrong (often last-iteration) value.

ℹ️ 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".

[duburcqa]

Addressed in 5498454. The Codex review correctly identified that UnaryOpType::tan was missing from NonLinearOps::unary_collections, but fixing only the set membership does not actually fix the bug it describes: a test with a = tan(a) in a dynamic loop still returns a wrong gradient because MakeAdjoint's tan formula was adj(stmt) * (1 + sqr(stmt)), and BackupSSA spills the forward tan result (stmt) to a plain alloca, so the reversed loop reads the last-iteration value for every backward step. The same latent bug also affected tanh and exp. The fix is therefore two-part: add tan to the non-linear set so the operand's alloca is promoted to an AdStack, and rewrite tan/tanh/exp's reverse formula to recompute the op on the adstack-backed operand rather than reuse the forward stmt. Covered by the new parametrized test_adstack_unary_loop_carried, which uses torch.autograd as the reference.

[duburcqa]

@claude review

[duburcqa]

@claude review

[claude]

LGTM — all previously flagged issues (tan/rsqrt missing from unary_collections, tan/tanh/exp forward-value reuse, alloca placement for GlobalLoadStmt outside the processed range-for, and the false-confidence test pattern) have been addressed in this revision.

Extended reasoning...

Overview

The PR fixes four orthogonal reverse-mode AD correctness bugs in quadrants/transforms/auto_diff.cpp: adjoint-alloca placement for GlobalLoadStmt outside the currently-processed RangeFor, a nested-if alloca placement bug, implementing tan derivatives, and fixing loop-carried primal recomputation for tan/tanh/exp. Two new test files (test_adstack.py, additions to test_ad_basics.py) cover the fixed paths.

Previous review round-trips

Three rounds of inline comments from this reviewer:

1. Round 1 — UnaryOpType::tan missing from NonLinearOps::unary_collections and tan/tanh/exp reverse formulas reusing the forward value rather than recomputing on the adstack-backed operand. Both addressed: tan added to the set; tan/tanh/exp now call tan(stmt->operand) / tanh(stmt->operand) / exp(stmt->operand).
2. Round 2 — UnaryOpType::rsqrt also missing from unary_collections. Addressed: rsqrt is now present alongside sqrt in the set.
3. Round 3 — Test comment claimed a purely-additive update was the sole promotion trigger, but a LocalLoad→LocalStore write-back cycle is independently sufficient. Addressed: test_adstack_unary_loop_carried now assigns a fresh from x[i] + j*0.1 inside the inner loop body, so LocalStore(a) precedes any LocalLoad(a) in the traversal order and the write-back cycle check cannot fire.

Security risks

None. This is an internal IR compiler pass with no user-visible API surface, no I/O, and no authentication/permission logic.

Level of scrutiny

Medium-high: correctness-critical AD transform logic, but the diff is focused, the reasoning in code comments is sound, and all identified gaps have been closed. The torch.autograd reference comparison in test_adstack_unary_loop_carried provides an independent oracle. No design decisions remain open.

Other factors

Bug hunting system found no new issues. The existing AD regression suite is documented as remaining green. Changes are backend-agnostic.

[claude]

---

Additional findings (outside current diff — PR may have been updated during review):

• 🟣 quadrants/transforms/auto_diff.cpp:1869-1876 — MakeDual::visit(UnaryOpStmt) has no handler for UnaryOpType::floor or UnaryOpType::ceil — the else-branch fires QD_NOT_IMPLEMENTED. This is a pre-existing gap: MakeAdjoint already has an explicit 'do nothing' branch for floor/ceil (piecewise-constant => zero gradient everywhere), but MakeDual does not. Any kernel using qd.floor() or qd.ceil() differentiated in forward mode (qd.ad.FwdMode) will crash at runtime with an unimplemented error. Since this PR directly modifies MakeDual::visit(UnaryOpStmt) to add the tan handler, this was the natural place to add the matching floor/ceil no-op.

  Extended reasoning...

  What the bug is and how it manifests

  MakeDual::visit(UnaryOpStmt) handles: neg, abs, sin, cos, tan (newly added by this PR), tanh, asin, acos, exp, log, sqrt, rsqrt, cast_value, logic_not — but has no branch for UnaryOpType::floor or UnaryOpType::ceil. The else-branch at the bottom of that function calls QD_NOT_IMPLEMENTED. Any kernel compiled in forward-mode AD (qd.ad.FwdMode) that contains a qd.floor() or qd.ceil() call will therefore crash at runtime with an unimplemented error.

  The specific code path that triggers it

  In auto_diff.cpp, MakeDual::visit(UnaryOpStmt): the chain of else-if branches checks each op type in order. After logic_not (which does nothing), the final else clause fires QD_NOT_IMPLEMENTED with the op type name. UnaryOpType::floor and UnaryOpType::ceil are not listed anywhere in MakeDual::visit(UnaryOpStmt), so both fall through to this crash.

  Why existing code does not prevent it

  MakeAdjoint::visit(UnaryOpStmt) was already written correctly. Its very first branch is:
  if (stmt->op_type == UnaryOpType::floor || stmt->op_type == UnaryOpType::ceil) { // do nothing }
  This is correct because floor and ceil are piecewise-constant: their derivative is 0 everywhere (except integer points where it is undefined). No accumulation is needed. MakeDual simply never received the corresponding no-op branch; the asymmetry predates this PR.

  What the impact would be

  Any kernel containing qd.floor() or qd.ceil() that is differentiated via qd.ad.FwdMode will hit QD_NOT_IMPLEMENTED and crash at runtime. The crash is a hard failure, not a silent correctness issue, but it is surprising because floor/ceil are supported ops and reverse-mode handles them without issue. The bug is exposed by the test_ad_frac test in test_ad_basics.py which uses qd.floor and qd.ceil, though that test only exercises reverse mode.

  How to fix it

  Add the floor/ceil no-op as the first branch in MakeDual::visit(UnaryOpStmt), mirroring MakeAdjoint:
  } else if (stmt->op_type == UnaryOpType::floor || stmt->op_type == UnaryOpType::ceil) {
  // do nothing: floor/ceil are piecewise-constant, dual component is zero

  Step-by-step proof

  1. User writes qd.floor(x[i]) inside a kernel, differentiates with qd.ad.FwdMode.
  2. The forward pass runs normally since UnaryOpType::floor computes fine.
  3. MakeDual::run(block) is invoked on the independent block.
  4. MakeDual::visit(UnaryOpStmt) is called for the floor statement.
  5. stmt->op_type == UnaryOpType::floor matches none of the existing branches (neg, abs, sin, cos, tan, tanh, asin, acos, exp, log, sqrt, rsqrt, cast_value, logic_not).
  6. The final else branch fires: QD_NOT_IMPLEMENTED and the program crashes.
  7. The same happens for ceil, which is likewise absent from MakeDual.

• 🟡 quadrants/transforms/auto_diff.cpp:1869-1873 — The comment at auto_diff.cpp line 1872 in MakeDual::visit(tan) says 'See the matching reverse-mode case for rationale,' but the reverse-mode rationale explains why it avoids reusing the forward stmt value (BackupSSA overwrites it each iteration), while MakeDual correctly does reuse stmt — the opposite design choice. A maintainer following the cross-reference may misread it as a reason to also recompute tan(stmt->operand) in forward mode, which would be redundant but reflects a misunderstanding of the design.

  Extended reasoning...

  What the bug is and how it manifests

  MakeDual::visit(tan) at auto_diff.cpp line 1872 uses sqr(stmt) — it reuses the forward tan result inline with the JVP. The comment says 'See the matching reverse-mode case for rationale.' But the reverse-mode comment's entire rationale explains why MakeAdjoint avoids using stmt: 'the primal is per-iteration inside dynamic loops but BackupSSA only spills forward values to a single plain alloca, so reading the forward tan would use the last-iteration value.' A reader following the cross-reference lands on an explanation for avoiding stmt, while the actual forward-mode code uses stmt — the opposite design choice with opposite reasoning.

  The specific code path that triggers it

  MakeDual::visit(UnaryOpStmt) for tan (auto_diff.cpp ~line 1872): accumulate(stmt, mul(add(constant(1), sqr(stmt)), dual(stmt->operand))). sqr(stmt) reuses the forward tan value, which is correct for forward mode because MakeDual generates JVPs inline during the forward pass — stmt is always the current-iteration value with no BackupSSA concern. The comment then tells the reader to look at MakeAdjoint::tan for rationale. MakeAdjoint::tan says it recomputes tan(operand) specifically because the forward stmt is unreliable due to BackupSSA. These rationales are for opposite choices.

  Why existing code does not prevent it

  The cross-reference is the only documentation for why MakeDual::tan uses stmt. No test validates comment accuracy. The consistent tanh case in MakeDual (sqr(stmt), no comment) shows the correct self-contained pattern.

  What the impact would be

  A future maintainer reads MakeDual::tan, follows the comment to MakeAdjoint, reads that stmt is unsafe due to BackupSSA, and may change MakeDual to recompute tan(stmt->operand). The change would be functionally redundant but reflects a misunderstanding. In the worst case the maintainer might also introduce a gratuitous consistency requirement between forward and reverse that leads to incorrect future changes.

  How to fix it

  Replace the cross-reference with a self-contained explanation: 'In forward mode the JVP is computed inline with the forward pass so stmt always holds the current-iteration tan value; BackupSSA does not apply here, unlike the reverse-mode case where the operand must be recomputed.'

  Step-by-step proof of the misleading cross-reference

  1. MakeDual::tan code: accumulate(stmt, mul(add(constant(1), sqr(stmt)), dual(stmt->operand))). Uses stmt directly — correct because MakeDual runs inline with the forward pass.
  2. Comment says: 'See the matching reverse-mode case for rationale.'
  3. Reader navigates to MakeAdjoint::tan. Its comment says: 'Recompute tan(operand) rather than reusing the forward value: the primal is per-iteration inside dynamic loops but BackupSSA only spills forward values to a single plain alloca, so reading the forward tan would use the last-iteration value in the reversed loop.'
  4. Reader infers: stmt is dangerous because BackupSSA overwrites it. MakeDual should also not reuse stmt.
  5. Conclusion 4 is wrong for forward mode — BackupSSA runs after MakeAdjoint on the reverse pass, not on the forward pass that MakeDual generates. The cross-reference is pointing to a rationale that does not apply.

• 🟡 tests/python/test_adstack.py:106-128 — The abs parametrization in test_adstack_unary_loop_carried (line 103) uses x_val=0.5, giving a=0.5, 0.6, 0.7 for j=0,1,2 — all positive — so the correct gradient (sgn(0.5)+sgn(0.6)+sgn(0.7)=3) and the broken gradient (3*sgn(0.7)=3) are identical; the test cannot detect if abs is accidentally dropped from NonLinearOps::unary_collections. This is a pre-existing test quality weakness unrelated to the core fixes in this PR. Using x_val=-0.15 (giving a=-0.15,-0.05,0.05, correct gradient=-1, broken gradient=3) would make the assertion actually sensitive to the regression the comment claims to guard against.

  Extended reasoning...

  What the bug is and how it manifests

  The test comment at test_adstack.py lines ~108-118 asserts: "If qd_op is dropped from [unary_collections], a_alloca stays a plain alloca ... producing a wrong gradient that torch.autograd catches." For the abs parametrization this claim is false. With x_val=0.5 and j=0,1,2, the kernel evaluates a = x[i] + j*0.1 giving a=0.5, 0.6, 0.7 — every value strictly positive. abs'(a) = sgn(a) = +1 for all three, so the summed gradient is unconditionally 3.

  The specific code path that triggers it

  The test kernel assigns a fresh each inner iteration (a = x[i] + j*0.1), so a_alloca is NOT loop-carried. The comment is correct that AdStackAllocaJudger::visit(LocalStoreStmt) write-back cycle check cannot fire independently, and the only promotion path is via AdStackAllocaJudger::visit(UnaryOpStmt) checking NonLinearOps::unary_collections membership. If abs were removed from that set, a_alloca would remain a plain AllocaStmt, BackupSSA would spill its last-written value (a=0.7) to a single plain alloca, and the reversed inner loop would read 0.7 for every backward step.

  Why existing code does not prevent it

  With the stale-value bug active: gradient = 3 * sgn(0.7) = 3 * 1 = 3. With the correct adstack promotion: gradient = sgn(0.5)+sgn(0.6)+sgn(0.7) = 1+1+1 = 3. Both evaluate to 3. The assert x.grad[0] == approx(x_t.grad.item()) comparison succeeds in both cases, so the regression goes undetected.

  What the impact would be

  This is a test-coverage weakness, not a production correctness bug. abs was already in unary_collections before this PR, and the production code remains correct. The risk is that a future accidental deletion of abs from unary_collections would pass this test silently, defeating the regression guard the comment advertises.

  How to fix it

  Replace x_val=0.5 with x_val=-0.15. This gives a=-0.15, -0.05, 0.05 for j=0,1,2. The correct gradient is sgn(-0.15)+sgn(-0.05)+sgn(0.05) = -1-1+1 = -1. The broken gradient (stale a=0.05) is 3sgn(0.05) = 31 = 3. The assertion -1 != 3 now actually catches the regression.

  Step-by-step proof

  1. x_val=0.5, inner loop j=0,1,2: a_0=0.5, a_1=0.6, a_2=0.7 — all positive.
  2. Hypothetically drop abs from unary_collections: AdStackAllocaJudger never sets is_stack_needed_=true for a_alloca. ReplaceLocalVarWithStacks leaves it as a plain AllocaStmt.
  3. BackupSSA: a_alloca is spilled to a single plain alloca. Forward pass overwrites it with a_0, then a_1, then a_2=0.7. After the full forward pass the plain alloca holds 0.7.
  4. Reversed inner loop reads 0.7 three times. Gradient = 3*sgn(0.7) = 3.
  5. torch.autograd also returns 3 (sgn(0.5)+sgn(0.6)+sgn(0.7)=3). Assertion: 3 == approx(3) — passes. Regression silently missed.

[hughperkins]

Checklist:

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

=> ok to merge

[claude]

LGTM — test-only PR, all substantive issues addressed, human reviewer has signed off.

Extended reasoning...

Overview

Single-file PR adding 182 lines of regression tests to tests/python/test_adstack.py. No production code changes. Adds test_adstack_basic_gradient, test_adstack_basic_gradient_f64, test_adstack_basic_gradient_negative, test_adstack_sum_linear, and test_adstack_sum_linear_f64 covering reverse-mode AD through dynamic and static for-loops with the adstack extension.

Security risks

None — test-only diff with no authentication, crypto, or permission-sensitive paths.

Level of scrutiny

This is new test code with no production impact. The merge risk is essentially zero. The changes are self-contained in the test file and follow established patterns in the test suite.

Other factors

The PR went through extensive iterative review (10+ rounds of inline comments) covering: missing unary_collections entries (fixed in sibling PRs #501/#503), redundant decorator kwargs, inaccurate docstrings, pytest.approx tolerance floor behaviour, dead parameters, and comment accuracy. All red and yellow items were either fixed directly or split into dedicated follow-up PRs. The one remaining documentation nit (the "f32-narrowing regression" justification comment) is inaccurate in its stated reason but the code it justifies — using abs_tol=0 with pytest.approx — is correct. The human reviewer saw this outstanding comment and explicitly said "ok to merge". The bug hunting system found no bugs in the current revision.