[AutoDiff] Autodiff 2: Implement derivative for tan

[duburcqa]

Implement derivative for qd.tan

Replaces QD_NOT_IMPLEMENTED in reverse and forward mode, registers tan as non-linear so its operand rides the adstack in dynamic loops, and pins the NaN-propagation open question with a strict xfail.

TL;DR

Reverse mode:

} else if (stmt->op_type == UnaryOpType::tan) {
  // d/dx tan(x) = 1 + tan(x)^2. Recompute tan(operand) rather than reuse the forward stmt:
  // BackupSSA spills forward values to a single plain alloca overwritten each iteration, so
  // reading the forward `stmt` in a reversed dynamic loop would use the last-iteration value.
  // `operand` rides the adstack through its LocalLoad, so a fresh tan on it is per-iteration correct.
  accumulate(stmt->operand,
             mul(adjoint(stmt),
                 add(constant(1, stmt->ret_type), sqr(tan(stmt->operand)))));

Forward mode uses the same identity but reuses stmt (safe because MakeDual runs in primal order — no BackupSSA stale-value concern).

Why

qd.tan was the only trig primitive that hit QD_NOT_IMPLEMENTED on differentiation. The old comment mentioned NaN concerns near π/2; that concern is real but orthogonal to the basic formula, and the formula itself is straightforward (sec²(x) = 1 + tan²(x)). The NaN-propagation semantics question is pinned separately as an xfail so a later design decision on reverse-mode NaN poisoning doesn't block implementing the common case.

Changes

quadrants/transforms/auto_diff.cpp

• NonLinearOps::unary_collections — add UnaryOpType::tan. Without this, AdStackAllocaJudger would not promote the operand alloca in dynamic loops, and the reverse pass would silently read the last-iteration operand at n_iter >= 3.
• ADTransform::tan(Stmt*) — new IR builder helper, mirrors the existing sin / cos helpers.
• MakeAdjoint::visit(UnaryOpStmt*) — new UnaryOpType::tan branch, accumulate(operand, adjoint * (1 + tan(operand)²)). Comment spells out the recompute-vs-reuse choice (so the next reviewer in the chain doesn't "fix" it back to sqr(stmt)).
• MakeDual::visit(UnaryOpStmt*) — new UnaryOpType::tan branch, accumulate(stmt, (1 + sqr(stmt)) * dual(operand)). Comment notes that forward mode runs in primal order so reusing stmt is safe.

tests/python/test_ad_basics.py

Adds (qd.tan, np.tan) to test_poly's parametrize — smoke coverage outside dynamic loops.

tests/python/test_adstack.py

• Adds ("tan", 0.05, 0.0) to _UNARY_OPS_PARAMS. Inherits the full x_val × n_iter matrix from Autodiff 1 in both f32 and f64.
• New test_adstack_nan_propagation — parametrized over (log, -0.3), marked xfail(strict=True). Pins the open question: PyTorch's backward pass poisons x.grad with NaN when the forward is NaN; Quadrants currently evaluates the analytical formula (e.g. 1 / operand = -3.33) and returns a finite number. Either behaviour is defensible; the xfail documents the current divergence so a deliberate change of semantics in either direction forces a reviewer decision. Only log is parametrized — sqrt / asin / acos would XPASS under strict=True because their reverse formulas divide by sqrt(<=0) which is itself NaN.

Side-effect audit

Concern	Verdict
Multi-iteration reverse-mode	Now correct for tan; pinned by the extended _UNARY_OPS_PARAMS.
Forward mode	Correct; pinned by test_ad_basics::test_poly + Autodiff 3's test_unary_forward_mode_derivative.
Near-singularity behaviour	Not addressed; bounded operand a_max ≈ 0.999 in the adstack test keeps us away from π/2 ≈ 1.57. If tan near-singularity gradient behaviour ever becomes a concern, a dedicated test can be added.
NaN propagation	Pinned as xfail(strict=True) — deliberate design-decision deferral.

Stack

Autodiff 2 of 13. Based on #500 (baseline). Followed by #502 (tanh/exp).

[claude]

LGTM — the derivative is mathematically correct, follows the established recompute-from-operand pattern for reverse mode, and is properly gated by .

Extended reasoning...

Overview

This PR adds autodiff support in (reverse and forward mode), adds a helper to , adds to , and extends both and with coverage.

Security risks

None. This is a pure numerical/compiler transform change with no I/O, authentication, or data-access implications.

Level of scrutiny

Low-to-medium. The change is small, self-contained, and follows an already-established pattern. The derivative formula is correct; the decision to recompute in reverse mode rather than reusing the forward stmt value is the right call for the same reason it is right for /. Adding to is the natural companion fix. Both new test vectors (static-loop via and dynamic-loop via ) are well-chosen.

Other factors

The bug reports flag two pre-existing issues (stale forward values for / in dynamic loops, and absent from ) and one nit about a cross-reference comment in . These are independent of the correctness of the implementation introduced here.

[hughperkins]

how are we handling the nans issue that the old comment alludes to?

[duburcqa]

I rewrote the MakeDual tan comment to be self-contained, so it no longer alludes to any NaN concern - the forward-mode rationale is just "stmt is the current-iteration tan value because MakeDual runs in primal order, so reusing it is per-iteration correct." The previous comment's cross-reference to the reverse-mode rationale was the source of the NaN implication (BackupSSA talks about spilled-value staleness, not NaNs), and it was misleading.

On a dedicated NaN/Inf unit test: I don't think it adds value in this PR. Reverse-mode NaN semantics across Quadrants vs PyTorch are genuinely divergent (forward log(-0.3) is NaN in both, but backward d log(a) / da = 1/a gives -3.33 in Quadrants vs NaN in torch because torch poisons the backward graph). That's a separate correctness question, orthogonal to the tan derivative and the unary_collections contract; landing a NaN test here would either bake in the current (possibly wrong) Quadrants behaviour or fail on arrival, neither of which is useful. Happy to file a follow-up to pin down the intended NaN/Inf behaviour across forward and reverse mode if you think that's worth it.

[hughperkins]

Lets add an xfail for that test then please. A gradient of -3.33 seems pretty wrong to me.

[hughperkins]

Opus summary:

Summary

Implements the derivative of tan for both reverse-mode and forward-mode AD, replacing
the previous QD_NOT_IMPLEMENTED stub. Also registers tan as a non-linear unary op
so that the AD system correctly saves per-iteration values inside dynamic loops.

Changes

• Reverse mode (MakeAdjoint): Computes d/dx tan(x) = 1 + tan(x)², recomputing
  tan(operand) from the operand rather than reusing the forward result. This avoids a
  bug where the backward pass inside a dynamic loop would read the last-iteration value
  instead of the correct per-iteration value.
• Forward mode (MakeDual): Same formula, 1 + tan(x)², using the forward tan
  result directly (safe in forward mode since values are consumed in-order).
• NonLinearOps::unary_collections: Adds UnaryOpType::tan so the AD system knows
  to save per-iteration copies of its operand in dynamic loops.
• ADTransform: Adds a tan() helper method used by the new derivative logic.
• Tests: Adds tan to both the existing test_ad_basics parametrized derivative
  test and the test_adstack dynamic-loop regression test.

Strengths

• Unblocks a previously stubbed-out op. tan was the only trig function that hit
  QD_NOT_IMPLEMENTED on differentiation; this brings it in line with sin/cos.
• Correct recompute strategy. The reverse-mode case deliberately recomputes
  tan(operand) instead of reusing the forward statement, which is the same pattern
  needed for correctness inside dynamic loops. The comment explains why.
• Good test coverage. Both the simple single-kernel AD test (test_ad_basics) and
  the dynamic-loop adstack regression test (test_adstack) are extended, covering the
  two distinct failure modes (wrong formula vs. wrong saved value).
• Small, focused diff. Easy to review — the C++ change is localised to the two
  visitor methods plus the non-linear op set.

Weaknesses / things to consider

• No test near singularities. tan diverges at x = π/2 + nπ. The existing tests
  use generic input values (0.5 in adstack, a small range in ad_basics) that stay well
  away from poles. A test that intentionally approaches a singularity would document the
  expected behaviour (large but finite gradient vs. inf/NaN).
• Forward mode reuses the forward tan statement directly (sqr(stmt)) while
  reverse mode recomputes from the operand. A brief comment in the forward-mode case
  explaining why reuse is safe there (no reversed loop) would make the asymmetry
  clearer.
• The 1 + tan² form can overflow earlier than 1/cos². For large |x| where
  tan(x) is large but finite, squaring it may overflow to inf in float32 while
  1/cos² would not. This is unlikely to matter in practice but worth noting.
• Comment in reverse mode still references internal concepts (BackupSSA,
  LocalLoad, adstack). These are meaningful to compiler developers but could
  benefit from a one-sentence plain-English lead-in for context.

[duburcqa]

@claude review

[claude]

LGTM — all prior review comments addressed; the spurious require=adstack on the NaN test is a nit already captured as an inline comment.

Extended reasoning...

Overview

This PR implements the reverse-mode and forward-mode derivative for qd.tan, replacing the QD_NOT_IMPLEMENTED stub. Changes span auto_diff.cpp (MakeAdjoint + MakeDual branches, tan() helper, NonLinearOps::unary_collections), test_ad_basics.py (adds qd.tan to the parametrized trig test), and test_adstack.py (adds qd.tan to the dynamic-loop regression test and adds a new NaN-semantics xfail test).

Security risks

None. This is a pure compiler transform / test change with no security surface.

Level of scrutiny

Low-to-medium. The diff is small and focused. The reverse-mode formula (1 + tan(x)^2) is standard; the recompute-from-operand pattern (to avoid BackupSSA stale values in dynamic loops) matches the documented intent and is consistent with the existing pattern for other ops. The forward-mode path correctly reuses stmt with a self-contained comment explaining why that is safe.

Other factors

All three issues from my prior review have been resolved or explicitly deferred to follow-up PRs (#502 for tanh/exp, #503 for rsqrt), which is a reasonable scope decision. The one new nit — test_adstack_nan_propagation carrying a spurious require=qd.extension.adstack guard — is captured as an inline comment and does not affect correctness on adstack-capable platforms. Test coverage is solid for both the static-loop and dynamic-loop paths.

[claude]

LGTM — all feedback addressed, tan derivative and test coverage look correct.

Extended reasoning...

Overview

PR adds the reverse-mode and forward-mode derivatives for qd.tan, replacing the old QD_NOT_IMPLEMENTED stub. Three files are touched: auto_diff.cpp (gradient logic + unary_collections membership), test_ad_basics.py (static-loop path), and test_adstack.py (dynamic-loop path + NaN xfail).

Security risks

None. This is a pure math/compiler transform change with no user-facing authentication, I/O, or memory-management surface.

Level of scrutiny

Moderate. The change is localised and follows the established recompute-from-operand pattern already present for sin/cos. The critical correctness property — recomputing tan(operand) in reverse mode rather than reusing the BackupSSA-spilled forward value — matches the pattern the PR description explicitly calls out and is verified by the dynamic-loop adstack test.

Other factors

All three of my prior inline comments (misleading cross-reference comment, tanh/exp stale-value pre-existing bug, rsqrt missing from unary_collections) are resolved: the comment was rewritten, the tanh/exp and rsqrt fixes are tracked in downstream PRs #502 and #503. The human reviewer's NaN xfail request was implemented and the spurious require=adstack guard I flagged was dropped. No outstanding unresolved threads remain.

[hughperkins]

Updated opus descripton:

Summary

Implements reverse- and forward-mode autodiff for qd.tan, replacing the previous QD_NOT_IMPLEMENTED stub. Adds tan to NonLinearOps::unary_collections so it goes
through the per-iteration adstack path in dynamic loops, wires it into the existing AD smoke tests, and uses the loop-carried regression test from the baseline PR as the
primary correctness vehicle.
Also adds a separate xfail(strict=True) test that pins an unresolved design question: how the reverse pass should handle NaN-producing forward inputs (PyTorch poisons the
gradient with NaN, Quadrants currently evaluates the analytical formula and returns a finite number).

What's in the PR

Net diff vs origin/duburcqa/split_adstack_unary_baseline: 3 files, +68 / -14.

quadrants/transforms/auto_diff.cpp

• Adds UnaryOpType::tan to NonLinearOps::unary_collections so the AD transform recognizes tan as a nonlinear unary op that needs per-iteration operand spilling on the
  adstack.
• Adds a tan IR builder helper to ADTransform.
• Reverse mode (MakeAdjoint): d/dx tan(x) = 1 + tan(x)^2, implemented as accumulate(operand, adjoint(stmt) * (1 + tan(operand)^2)). Comment justifies recomputing
  tan(operand) rather than reusing the forward stmt: BackupSSA only spills forward values to a single plain alloca, so reusing the forward tan would read the
  last-iteration value in a reversed dynamic loop. The operand, in contrast, rides the adstack via its LocalLoad, so a fresh tan on it is per-iteration correct.
• Forward mode (MakeDual): accumulate(stmt, (1 + tan(stmt)^2) * dual(operand)). Comment notes that forward mode runs in primal order, so reusing stmt is safe (no
  BackupSSA stale-value concern).

tests/python/test_ad_basics.py

• Adds (qd.tan, np.tan) to the existing scalar test_poly AD parametrize list — basic smoke coverage outside dynamic loops.

tests/python/test_adstack.py

• Adds ("tan", 0.05, 0.0) to the loop-carried unary parametrize introduced in the baseline PR. Comment notes that tan's singularity at π/2 ≈ 1.57 lies safely outside the
  positive-path a_max = 0.999.
• Adds a new test_adstack_nan_propagation test, parametrized over [("log", -0.3)] and marked xfail(strict=True). Pins the open design question about reverse-mode NaN
  handling. Docstring explicitly notes why sqrt/asin/acos were removed from the parametrize: their reverse formulas themselves divide by sqrt(<=0) which is NaN, so
  Quadrants actually does return NaN for those — which would XPASS under strict=True.

Good points

• Closes a known gap identified by the baseline regression test: tan was the most obvious unary op missing from NonLinearOps::unary_collections, and this PR fixes it
  on both forward and reverse paths.
• Implements both AD modes, not just reverse — keeps MakeAdjoint and MakeDual consistent.
• Comments document the non-obvious choice in the reverse formula (recompute tan(operand) rather than reuse stmt) and explain why it differs from forward mode. This
  is exactly the kind of footgun that the baseline test was designed to catch, and the comment makes the trap explicit for future contributors.
• Tests at three levels: smoke test in test_ad_basics.py (scalar), loop-carried test in test_adstack.py (dynamic loop, multi-iteration), and the new NaN-propagation
  pin. Good coverage triangle.
• NaN test is honest about scope. It uses xfail(strict=True) to document current behavior as an unresolved design question rather than silently pick a side. The
  docstring explains why only log is parametrized and what would happen for the sibling ops — leaves a clear trail for whoever picks up the semantics decision.
• Replaces QD_NOT_IMPLEMENTED with a working implementation, so user code using qd.tan under autodiff stops hitting a hard error.

Bad points / concerns

• (1 + tan(x)^2) vs 1 / cos(x)^2. Mathematically equivalent, but the prior code's comment specifically called out singularity-induced NaNs as the reason for not
  implementing the derivative. The new formula avoids the explicit 1/cos^2 division but tan(x)^2 still blows up near π/2, so the original concern about hard-to-debug NaNs
  hasn't really been addressed — just relocated. Worth a brief note in the PR body or commit message acknowledging this trade-off (and the xfail NaN test arguably leans into
  it rather than resolving it).
• Reverse-mode tan is recomputed twice per backward step: once for tan(operand)^2 and any future use. The comment justifies why reuse is unsafe under BackupSSA, but
  it's worth noting that the per-iteration operand spill on the adstack plus the recomputation costs one extra tan evaluation in the reversed loop. Probably fine, but a
  perf-sensitive reviewer may want benchmarks.
• Single xfail NaN case. The test only parametrizes log because sqrt/asin/acos happen to XPASS under strict=True. That asymmetry is worth surfacing in a
  follow-up: either pick the semantics and update all four uniformly, or document the divergent behavior explicitly. As it stands, the xfail pins one half of an inconsistent
  story.
• No coverage of tan near its singularity. The loop-carried test is deliberately bounded to the positive-path a_max ≈ 0.999, well below π/2. Reasonable for a
  regression pin (we're testing AD plumbing, not numerical robustness near singularities), but it means a future change that breaks tan's gradient near π/2 would slip
  through. Probably acceptable for this PR; worth a follow-up if tan derivative behavior near singularities ever becomes a user-visible concern.
• No f64 variant of either the new tan parametrize entry or the NaN test. Inherits the same f32-only limitation as the baseline test.
• Merge commit, not rebase. Branch was brought up to date via git merge rather than rebased onto the baseline. History contains a merge commit; depending on repo
  conventions a clean rebase + force-push may be preferred.
• Two PRs in flight on the same surface. This PR depends on duburcqa/split_adstack_unary_baseline (the new test parametrize entry for tan only makes sense once that
  test exists). Reviewers will need to merge them in order, and a force-push or rebase on the baseline PR will require re-syncing this one.

[hughperkins]

I'm not sure I follow. Why cant we only xfail log, and not xfail the others?

[hughperkins]

please parmetrize over fp32 vs fp64

[claude]

LGTM — all prior review comments addressed; the tan derivative implementation, adstack registration, and NaN-semantics xfail test are all correct.

Extended reasoning...

Overview

This PR adds the reverse-mode and forward-mode derivative for tan, replaces the QD_NOT_IMPLEMENTED stub in both MakeAdjoint and MakeDual, registers UnaryOpType::tan in NonLinearOps::unary_collections, and extends both test_ad_basics and test_adstack with tan coverage. A new test_adstack_nan_propagation xfail test documents the NaN-semantics divergence between Quadrants and PyTorch.

Security risks

None. This is purely a compiler transform and test change with no user-facing data handling.

Level of scrutiny

Low-to-medium. The change is small and self-contained: the derivative formula 1 + tan(x)^2 is mathematically straightforward, the recompute-from-operand pattern (to avoid BackupSSA stale-value) follows the same convention already established in the codebase, and both the static-loop and dynamic-loop paths are tested. The core correctness concern (per-iteration operand recompute vs. stale forward value) is well-documented in the comments.

Other factors

All prior inline comments from my previous review run have been resolved. The remaining open item — a docstring nit about why tanh/exp are absent from the parametrize list — is minor and already captured as an inline comment; it does not affect correctness. The xfail decorator change (dropping the spurious require=qd.extension.adstack guard) was also correctly applied. No outstanding reviewer threads remain unaddressed.

[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

[claude]

LGTM — tan derivative is correctly implemented in both reverse and forward mode, adstack registration is in place, and all prior review concerns have been addressed.

Extended reasoning...

Overview

The PR adds tan autodiff support across three files: the C++ AD transform (MakeAdjoint + MakeDual + NonLinearOps::unary_collections), a smoke test entry in test_ad_basics.py, and dynamic-loop regression coverage plus NaN-semantics xfail tests in test_adstack.py. The diff is focused and self-contained.

Security risks

None. This is pure compiler/AD transform logic with no I/O, authentication, or privilege boundaries involved.

Level of scrutiny

Moderate. The core formula (d/dx tan(x) = 1 + tan(x)^2) is straightforward and the recompute-vs-reuse choice in the reverse pass is correctly reasoned and documented. The test coverage triangle (scalar smoke, multi-iteration dynamic loop in f32/f64, NaN-semantics xfail) is solid. Prior review rounds addressed comment clarity, adstack guard misuse, and f64 coverage gaps.

Other factors

The human reviewer explicitly signed off. All previous bot-flagged issues are resolved or deferred to follow-up PRs with acknowledgment (#502 for tanh/exp recompute, #503 for rsqrt). The remaining nit — tan missing from test_ad_grad_check.py's numerical gradient check — is a test coverage gap, not a correctness defect in the implementation itself, and does not block merging.