[AutoDiff] Autodiff 3: Recompute tanh/exp on the operand in the reverse pass

[duburcqa]

Recompute tanh / exp on the operand in the reverse pass

Fixes silently-wrong gradients at n_iter >= 3 in dynamic loops by recomputing tanh(operand) / exp(operand) instead of reusing the forward UnaryOpStmt value. Same root cause and same fix pattern as Autodiff 2 applied to two more ops.

TL;DR

Reverse mode, before / after:

 } else if (stmt->op_type == UnaryOpType::tanh) {
-  accumulate(stmt->operand, mul(adjoint(stmt), sub(constant(1), sqr(stmt))));
+  // Recompute tanh(operand) in the reverse pass: BackupSSA spills the forward `stmt` to a
+  // single plain alloca overwritten each iteration, so reading it from a reversed dynamic loop
+  // would use the last-iteration value. `operand` rides the adstack, so a fresh tanh on it is
+  // per-iteration correct. Trade-off: tanh is evaluated twice per iteration (forward + backward);
+  // caching the forward value on the adstack is a future optimization.
+  accumulate(stmt->operand, mul(adjoint(stmt), sub(constant(1), sqr(tanh(stmt->operand)))));
 } else if (stmt->op_type == UnaryOpType::exp) {
-  accumulate(stmt->operand, mul(adjoint(stmt), stmt));
+  // See the tanh case above.
+  accumulate(stmt->operand, mul(adjoint(stmt), exp(stmt->operand)));

Why

Both ops were already in NonLinearOps::unary_collections, so their operand allocas get promoted to AdStackAllocaStmt in dynamic loops. But the reverse-mode formulas read the forward stmt:

• tanh: 1 - tanh(x)², where tanh(x) was substituted with the forward stmt.
• exp: exp(x), where exp(x) was substituted with the forward stmt.

BackupSSA spills each forward UnaryOpStmt value to a single plain alloca for later reverse-pass reads. Inside a dynamic loop that alloca gets overwritten every forward iteration, so by the time the reversed loop reads it, it holds the last-iteration value regardless of which reverse iteration is running. Gradients come out off by a factor proportional to the op's non-linearity across the iteration range.

Recomputing from the operand (which rides the adstack through its LocalLoad) makes the reverse pass per-iteration correct.

Changes

quadrants/transforms/auto_diff.cpp

• ADTransform::tanh(Stmt*) and ADTransform::exp(Stmt*) — new IR builder helpers, mirroring the existing ADTransform::tan helper from Autodiff 2.
• MakeAdjoint::visit(UnaryOpStmt*):
  • UnaryOpType::tanh branch — sqr(stmt) → sqr(tanh(stmt->operand)).
  • UnaryOpType::exp branch — stmt → exp(stmt->operand).
  • UnaryOpType::log and UnaryOpType::sqrt branches — no code change. New comments spell out explicitly that these do not need the same recompute: their reverse formulas already read stmt->operand directly (not the forward stmt), so BackupSSA is not in the critical path. Comments mirror the log block's phrasing so a future reviewer can quickly distinguish "already operand-based" from "needs the recompute workaround".

tests/python/test_adstack.py

• _UNARY_OPS_PARAMS: add ("tanh", 0.3, -0.4) and ("exp", 0.3, -0.4) to the real-domain group. The parametrize-level comment is rewritten to distinguish the sign-crossing requirement for abs / sin / cos from the mere-domain requirement for tanh / exp (both groups share the same (step, offset) because their domains are all reals; tanh / exp don't need the zero-crossing but it doesn't hurt them).
• New test_unary_forward_mode_derivative — pins MakeDual (forward mode) for tan / tanh / exp plus the audit-adjacent log / sqrt. Forward mode is safe-by-primal-order (no BackupSSA stale-value concern) but that invariant wasn't tested anywhere; now it is.

Side-effect audit

Concern	Verdict
Multi-iteration reverse-mode for tanh / exp	Fixed; regression pinned by _UNARY_OPS_PARAMS.
log / sqrt correctness	Unchanged (both were already operand-based). Comments make this explicit.
Forward mode	Unchanged; pinned by new test_unary_forward_mode_derivative.
Recompute cost	tanh / exp now evaluated twice per iteration. For exp in hot ML kernels this may be measurable; future optimization would cache the forward value on the adstack.

Stack

Autodiff 3 of 13. Based on #501 (tan). Followed by #503 (rsqrt).

[claude]

LGTM — the tanh and exp adjoint recomputes are mathematically correct and consistent with the established pattern for tan.

Extended reasoning...

Overview

This PR modifies two files: quadrants/transforms/auto_diff.cpp (adjoint formulas for tanh and exp in MakeAdjoint::visit(UnaryOpStmt*), plus two new helper methods on ADTransform) and tests/python/test_adstack.py (two new parametrize entries). The fix changes how the reverse pass computes tanh and exp gradients inside dynamic loops.

Security risks

None. This is a pure algorithmic/numerical correctness fix in the AD compiler pass. No auth, crypto, or permissions code is touched.

Level of scrutiny

Low-to-medium. The change is narrow and mechanical: two one-line formula substitutions and two helper stubs that mirror the existing tan() helper added in the predecessor commit. The math is straightforward — d/dx tanh(x) = 1 - tanh²(x) and d/dx exp(x) = exp(x) — and recomputing on the adstack-backed operand (rather than reusing the forward stmt value) is exactly the same pattern applied to tan in the base commit. Both tanh and exp are already present in NonLinearOps::unary_collections, so the operand alloca promotion path is already exercised.

Other factors

The test extension is well-designed: it uses PyTorch as a reference oracle, exercises a loop-carried operand pattern that specifically stresses the adstack promotion path, and the comment in the test body documents the precise invariant being validated. The pre-existing rsqrt gap (flagged by the bug hunter) is explicitly acknowledged in that comment and is orthogonal to this PR's correctness.

[hughperkins]

Opus review:

Summary

Fixes reverse-mode AD correctness for tanh and exp inside dynamic loops. The
existing derivative formulas were mathematically correct but reused the forward
statement value (stmt), which in a dynamic loop gets spilled to a single scratch slot
that every forward iteration overwrites. The reversed loop then reads the
last-iteration value for every backward step, producing silently wrong gradients.
The fix follows the same pattern already applied to tan in the parent PR: recompute
the op from the operand (which is correctly backed by the adstack) instead of reusing
the forward result.

Changes

• tanh reverse mode: sqr(stmt) → sqr(tanh(stmt->operand)).
• exp reverse mode: stmt → exp(stmt->operand).
• ADTransform: Adds tanh() and exp() helper methods for the recompute calls.
• Tests: Adds tanh and exp to the test_adstack_unary_loop_carried parametrize
  list.
• No forward-mode changes — forward mode executes in primal order so reusing stmt
  is already correct there.

Strengths

• Targeted, minimal fix. Each change is a single-expression substitution with a
  clear mechanical pattern (stmt → op(stmt->operand)), making it easy to review and
  verify.
• Consistent with the tan precedent. Uses the same recompute-from-operand strategy
  and the same comment structure established in the parent PR.
• Covered by existing parametrized tests. Adding the two ops to
  test_adstack_unary_loop_carried automatically inherits the multi-value, multi-field,
  multi-iteration cross-validation against PyTorch.
• No forward-mode churn. Correctly identifies that forward mode doesn't need the
  same treatment and leaves it untouched.

Weaknesses / things to consider

• Recompute cost. Both tanh and exp are now evaluated twice (once forward, once
  in the backward pass). For exp in particular this could be measurable in
  compute-heavy kernels. The previous code was wrong, so this is necessary for
  correctness, but a future optimisation could cache the recomputed value if the same
  operand feeds multiple uses.
• log and sqrt reverse cases still reuse stmt->operand directly (via
  div(adjoint, stmt->operand) and div(adjoint, mul(2, stmt)) respectively). These
  happen to be correct because their derivatives don't reference the op's own output,
  but a comment noting why they don't need the same recompute treatment would help
  future readers avoid accidentally "fixing" them.
• Comments still reference internal compiler concepts (BackupSSA, LocalLoad,
  adstack). Consistent with the parent PR's style, but a one-sentence plain-English
  lead-in would help readers unfamiliar with the compiler internals.
• sqrt reverse mode does reference stmt (the forward result) via
  div(adjoint, mul(2, stmt)). This looks like it could suffer the same stale-value
  bug — worth verifying that it's covered or flagging for a follow-up.

[hughperkins]

Could we add a short 1-sentence comment for each of:

• Recompute cost. Both tanh and exp are now evaluated twice (once forward, once
  in the backward pass). For exp in particular this could be measurable in
  compute-heavy kernels. The previous code was wrong, so this is necessary for
  correctness, but a future optimisation could cache the recomputed value if the same
  operand feeds multiple uses.
• log and sqrt reverse cases still reuse stmt->operand directly (via
  div(adjoint, stmt->operand) and div(adjoint, mul(2, stmt)) respectively). These
  happen to be correct because their derivatives don't reference the op's own output,
  but a comment noting why they don't need the same recompute treatment would help
  future readers avoid accidentally "fixing" them.
• sqrt reverse mode does reference stmt (the forward result) via
  div(adjoint, mul(2, stmt)). This looks like it could suffer the same stale-value
  bug — worth verifying that it's covered or flagging for a follow-up.

[duburcqa]

@claude review

[claude]

LGTM — comments from the prior review round are all addressed in the diff.

Extended reasoning...

Overview

The PR touches two files: quadrants/transforms/auto_diff.cpp (two new ADTransform helper methods and two one-expression substitutions in MakeAdjoint::visit(UnaryOpStmt*)) and tests/python/test_adstack.py (two new parametrize entries). The scope is tightly bounded to fixing the reverse-mode AD formulas for tanh and exp inside dynamic loops.

Security risks

None. This is a compiler IR transform for automatic differentiation. There are no auth, crypto, permission, or external-surface changes.

Level of scrutiny

Low. Each substantive change is a single-expression substitution (sqr(stmt) → sqr(tanh(stmt->operand)) for tanh; stmt → exp(stmt->operand) for exp) following the identical pattern established in the parent PR for tan. The fix is mechanically verifiable against the derivative identity and is covered by the extended parametrized test.

Other factors

The human reviewer's comment requests (recompute-cost note, log/sqrt rationale, sqrt stale-value clarification) are already incorporated into the diff as inline comments. The rsqrt pre-existing gap I raised was acknowledged and deferred to PR #503 with the inline thread resolved. No outstanding unresolved reviewer threads remain.

[claude]

LGTM — the tanh/exp recompute fixes are correct and minimal, the requested clarifying comments are present, and the rsqrt gap is tracked in #503.

Extended reasoning...

Overview

The PR modifies two files: auto_diff.cpp (adjoint formulas for tanh and exp in MakeAdjoint, plus two new ADTransform helper methods) and test_adstack.py (two new parametrize entries). The change is focused and follows the exact pattern established for tan in the parent PR.

Security risks

None. This is a compiler IR transformation with no auth, crypto, or permission logic.

Level of scrutiny

Low-to-medium. The fix is a one-expression substitution per op (sqr(stmt) → sqr(tanh(stmt->operand)) and stmt → exp(stmt->operand))) with a clear mechanical rationale. The test additions directly verify the corrected behavior against PyTorch autograd at multiple iteration counts.

Other factors

All substantive reviewer concerns have been resolved: the rsqrt gap is acknowledged and deferred to PR #503 (inline comment resolved), the clarifying comments requested by hughperkins (recompute cost, log/sqrt rationale, sqrt stale-value note) are present in the diff, and the remaining bug report is a documentation nit about the test parametrize comment that does not affect correctness or test coverage.

[claude]

LGTM — core fix is correct, all previous feedback addressed, remaining inline item is a comment-phrasing nit on a correct code path.

Extended reasoning...

Overview

This PR fixes reverse-mode AD for tanh and exp inside dynamic loops by replacing stale forward-value reuse (sqr(stmt) / stmt) with per-iteration recompute on the adstack-backed operand (sqr(tanh(stmt->operand)) / exp(stmt->operand)). It also adds tanh() and exp() helper methods to ADTransform, adds both ops to the test_adstack_unary_loop_carried parametrize, and clarifies comments on the log/sqrt/tanh/exp reverse cases. Two files changed: quadrants/transforms/auto_diff.cpp and tests/python/test_adstack.py.

Security risks

None. This is a compiler transform pass with no auth, crypto, or permission-sensitive code paths.

Level of scrutiny

Medium — this touches correctness-critical AD code, but the change is mechanical and follows an established pattern (identical fix was applied to tan in the parent PR #501). The parametrized tests cross-check against PyTorch autograd at multiple iteration counts and field values, giving strong regression coverage.

Other factors

All previously raised issues have been resolved: the rsqrt gap is acknowledged and deferred to PR #503; the parametrize comment contradiction I flagged has been corrected in the current diff to cleanly distinguish sign-crossing rationale (abs/sin/cos) from domain-only rationale (exp/tanh); and the comment improvements requested by the human reviewer are present in the diff. The sole remaining inline item is a nit about the sqrt backward comment drawing a misleading parallel to tanh/exp — the bug report itself notes the code is entirely correct and the risk is maintenance-documentation only.

[claude]

LGTM — all prior feedback addressed; the tanh/exp recompute fix is correct and the updated comments accurately explain the BackupSSA rationale.

Extended reasoning...

Overview

Two-file change: fixes the reverse-mode adjoint for and by recomputing the op on the adstack-backed operand instead of reusing the forward stmt value; adds and to the parametrized loop-carried gradient test.

Security risks

None. This is a pure compiler IR transformation pass with no I/O, auth, or data-exposure surface.

Level of scrutiny

Medium. The change touches AD correctness logic, but each modification is a single-expression substitution following the pattern established for in the parent PR. The math is straightforward (d/dx tanh(x) = 1 − tanh²(x), d/dx exp(x) = exp(x)), and the PyTorch cross-validation in the test directly catches wrong gradients.

Other factors

All three inline comments I raised in earlier rounds are resolved: (1) rsqrt missing from unary_collections is deferred to PR #503 with the author's acknowledgement; (2) the parametrize comment now correctly distinguishes the sign-crossing rationale (abs/sin/cos) from the domain-only rationale (exp/tanh); (3) the sqrt comment now mirrors the log framing and explicitly notes sqrt was already operand-based before this fix. The human reviewer's requested comments on recompute cost and log/sqrt correctness rationale are also present in the final diff.

[hughperkins]

Updated opus description:

Summary

Fixes the same BackupSSA stale-value bug for tanh and exp that the previous PR (split_autodiff_tan_derivative) fixed for tan: their reverse-mode formulas were
reading the forward stmt value, which in dynamic loops gets spilled to a single plain alloca overwritten each iteration. The reversed loop therefore read the last-iteration
value for every backward step, silently producing wrong gradients at n_iter ≥ 3.
Fix: recompute tanh(operand) / exp(operand) in the reverse pass against the adstack-backed operand instead of reusing the forward stmt. Adds the two ops to the
loop-carried regression parametrize, and documents log/sqrt as already-correct cases that don't need the workaround.

What's in the PR

Net diff vs origin/duburcqa/split_autodiff_tan_derivative: 2 files, +33 / -5.

quadrants/transforms/auto_diff.cpp

• Adds tanh and exp IR builders to ADTransform, mirroring the tan builder added in the previous PR.
• MakeAdjoint::tanh: 1 - sqr(stmt) (reused forward value) → 1 - sqr(tanh(stmt->operand)) (recompute on the operand).
• MakeAdjoint::exp: adjoint * stmt (reused forward value) → adjoint * exp(stmt->operand) (recompute on the operand).
• MakeAdjoint::log and MakeAdjoint::sqrt: no code change, but adds explanatory comments noting these reverse formulas already read stmt->operand directly (not the
  forward stmt), so they don't need the workaround. The sqrt comment explicitly notes it was already operand-based before this PR.
• MakeDual (forward mode): unchanged, consistent with the tan PR's reasoning that forward mode runs in primal order so reusing stmt is safe.

tests/python/test_adstack.py

• Adds ("tanh", 0.3, -0.4) and ("exp", 0.3, -0.4) to the loop-carried unary parametrize, in the "all-real-domain" group with sin/cos/abs.
• Rewrites the section comment to split the rationale: sin/cos/abs need sign-crossing operands (so their per-iteration gradient sign varies, defeating trivial
  pass-through); tanh/exp don't need crossing (their derivatives stay positive) but share the same (step, offset) because their domains are all reals.

Good points

• Closes two more instances of a now-known bug class. The tan PR established the pattern; this PR applies it systematically to tanh and exp. The fix is mechanical
  and easy to review against the previous one.
• Audits the rest of the unary table. Adds explicit "no workaround needed" comments to log and sqrt justifying why they're not affected (their reverse formulas read
  the operand directly, not the forward stmt). Closes the loop — a reviewer doesn't have to wonder whether other ops are still broken.
• Comments are honest about the trade-off. Explicitly notes that tanh/exp are now evaluated twice per iteration (once forward, once backward) and flags "caching the
  forward value on the adstack" as a future optimization. Doesn't oversell the fix.
• Test parametrize additions land naturally in the existing structure from the baseline PR. The rewritten section comment correctly distinguishes "needs sign-crossing for
  the test to be non-trivial" (abs/sin/cos) from "shares parameters incidentally" (tanh/exp), so a future reader doesn't infer a constraint that doesn't exist.
• Tightly scoped. Just the two op fixes plus comments and parametrize entries. No drive-by changes.

Bad points / concerns

• Performance regression for tanh/exp in tight backward loops. Each iteration now evaluates the op twice. For exp in particular this can be a meaningful cost in
  ML-style code. The comment acknowledges it and points at adstack caching as the future fix, but no benchmark is included to quantify the impact. Worth measuring before merge
  if any user kernel is known to be hot on exp.grad().
• Same root cause keeps getting patched op-by-op. This is the third PR (tan, tanh, exp after the baseline test) that fixes the same
  BackupSSA-spills-to-single-alloca pattern in essentially the same way. The repeated occurrences argue for a structural fix — e.g. either making BackupSSA use the adstack
  for ops in unary_collections, or making the reverse pass default to operand-recompute for any nonlinear unary — rather than continuing to audit each op individually. As-is,
  any new differentiable unary op added to the codebase needs a manual review of "does its reverse formula read stmt or stmt->operand?"
• No coverage that the audit on log/sqrt is correct. The "no workaround needed" claim is asserted in comments but not pinned by a test that would fail if a refactor
  accidentally regressed those formulas to read stmt. The existing parametrize would catch it (since log/sqrt are already in the test), so the coverage is de facto
  there, but the linkage is implicit.
• MakeDual (forward mode) untouched, untested for the new ops. Forward mode is argued safe-by-construction (primal order), but the regression test only exercises reverse
  mode. If forward mode ever stops running in strict primal order, the silent-NaN class of bug returns. Probably fine for this PR; worth a sentence in any future "AD
  architecture" doc.

[hughperkins]

Thoughts on:

• Same root cause keeps getting patched op-by-op. This is the third PR (tan, tanh, exp after the baseline test) that fixes the same
  BackupSSA-spills-to-single-alloca pattern in essentially the same way. The repeated occurrences argue for a structural fix — e.g. either making BackupSSA use the adstack
  for ops in unary_collections, or making the reverse pass default to operand-recompute for any nonlinear unary — rather than continuing to audit each op individually. As-is,
  any new differentiable unary op added to the codebase needs a manual review of "does its reverse formula read stmt or stmt->operand?"

?

Please address:

• No coverage that the audit on log/sqrt is correct. The "no workaround needed" claim is asserted in comments but not pinned by a test that would fail if a refactor
  accidentally regressed those formulas to read stmt. The existing parametrize would catch it (since log/sqrt are already in the test), so the coverage is de facto
  there, but the linkage is implicit.
• MakeDual (forward mode) untouched, untested for the new ops. Forward mode is argued safe-by-construction (primal order), but the regression test only exercises reverse
  mode. If forward mode ever stops running in strict primal order, the silent-NaN class of bug returns. Probably fine for this PR; worth a sentence in any future "AD
  architecture" doc.

[claude]

LGTM — all prior feedback addressed, fix is correct and well-covered by the parametrized tests.

Extended reasoning...

Overview

Modifies MakeAdjoint::visit(UnaryOpStmt*) for tanh and exp in auto_diff.cpp, adds matching tanh()/exp() helper methods to ADTransform, and extends test_adstack_unary_loop_carried in test_adstack.py with the two new ops.

Security risks

None. This is an internal compiler IR transform with no auth, crypto, or I/O surface.

Level of scrutiny

Low-to-medium. The fix is a one-expression substitution in two cases, following the identical pattern already established for tan in the parent PR (#501). The derivative formulas are mathematically straightforward and the operand-recompute strategy is well-understood in this codebase.

Other factors

All three issues from my prior reviews are resolved: the rsqrt gap is deferred to PR #503, the parametrize comment now correctly scopes the sign-crossing rationale to abs/sin/cos only, and the sqrt comment mirrors the log framing rather than drawing a false parallel to tanh/exp. Test coverage via PyTorch cross-validation catches wrong gradients at n_iter >= 3, which is exactly the regression path this fix closes.

[claude]

LGTM — all prior review concerns addressed.

Extended reasoning...

Overview

This PR fixes reverse-mode AD for tanh and exp inside dynamic loops by replacing the stale forward-stmt reuse with tanh(stmt->operand) / exp(stmt->operand) respectively. It also adds tanh() and exp() IR-builder helpers to ADTransform, extends the test_adstack_unary_loop_carried parametrize with the two new ops, adds a new test_unary_forward_mode_derivative test pinning forward mode for tan/tanh/exp/log/sqrt, and rewrites explanatory comments for log/sqrt/tanh/exp.

Security risks

None. This is a compiler transform for automatic differentiation with no security surface.

Level of scrutiny

Moderate — touches AD core logic, but the fix is a mechanical single-expression substitution following the exact pattern established by the tan predecessor PR. The parametrized PyTorch cross-check at n_iter=1/3/10 provides strong numerical validation, and the new forward-mode test pins the safety argument for MakeDual.

Other factors

All three concerns raised in my previous reviews have been addressed: the rsqrt gap is deferred to PR #503, the test comment contradiction was rewritten, and the misleading sqrt framing was corrected. The hughperkins requests for forward-mode coverage and log/sqrt audit linkage are addressed by the new test_unary_forward_mode_derivative. The one remaining inline nit (missing safety comment on MakeDual::tanh/exp mirroring MakeDual::tan) is documentation-only and has been filed as an inline comment.

[hughperkins]

Checklist:

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

=> ok to merge