Skip to main content

pounce_nl/
nl_tape.rs

1//! Flat-tape reverse-mode AD for `.nl` expression trees.
2//!
3//! Replaces the FD-based Hessian path with a port of the tape AD used
4//! in `ripopt::nl::autodiff`. The tape is a `Vec<TapeOp>` where each op
5//! refers to its operands by tape-slot index; forward evaluation runs
6//! through the slice once filling a parallel `Vec<f64>` of values, and
7//! reverse-mode adjoints walk the same buffer backwards.
8//!
9//! Sparse Hessians are computed by forward-over-reverse: for each
10//! variable `j` that the tape depends on, run a forward tangent sweep
11//! seeded with `e_j`, then a second-order reverse sweep that produces
12//! column `j` of the Hessian. The caller supplies a `(row, col) -> nnz
13//! position` map (lower triangle, row >= col), and contributions are
14//! accumulated in place — the outer loop in `eval_h` calls the same
15//! map for the objective and every active constraint, so every
16//! Lagrangian term lands in the right slot.
17//!
18//! Common subexpressions are tape-emitted **once**: when the recursive
19//! builder hits `Expr::Cse(rc)` it keys on the `Arc` pointer identity,
20//! emitting the body the first time and returning the cached
21//! result-slot index on subsequent references. The forward pass then
22//! computes each CSE once and the reverse pass folds adjoints from
23//! every reference into a single slot — exact chain-rule behaviour.
24
25use std::collections::{BTreeSet, HashMap, HashSet};
26use std::sync::Arc;
27
28use super::nl_external::{EvalResult, ExternalArg, ExternalLibrary, ExternalResolver};
29use super::nl_reader::{BinOp, CmpOp, Expr, FuncallArg, UnaryOp};
30
31/// One operation in the flattened tape. Operand fields are tape-slot
32/// indices into the same tape; `Var(i)` references problem variable
33/// index `i` (read from the input `x` slice during forward).
34#[derive(Debug, Clone)]
35pub enum TapeOp {
36    Const(f64),
37    Var(usize),
38    Add(usize, usize),
39    Sub(usize, usize),
40    Mul(usize, usize),
41    Div(usize, usize),
42    Pow(usize, usize),
43    Neg(usize),
44    Abs(usize),
45    Sqrt(usize),
46    Exp(usize),
47    Log(usize),
48    Log10(usize),
49    Sin(usize),
50    Cos(usize),
51    Tan(usize),
52    Atan(usize),
53    Acos(usize),
54    Sinh(usize),
55    Cosh(usize),
56    Tanh(usize),
57    Asin(usize),
58    Acosh(usize),
59    Asinh(usize),
60    Atanh(usize),
61    /// Gauss error function `erf(vals[a])` (issue #469). Unlike the other
62    /// transcendentals here it has no AMPL `.nl` opcode — AMPL has no `erf`
63    /// — so it is reachable only through the in-memory `Expr` path
64    /// (`UnaryOp::Erf`), which is exactly the case it was added for:
65    /// modeling frontends that build a tape directly instead of round-
66    /// tripping through `.nl`. Smooth everywhere, with closed-form
67    /// derivatives `erf'(u) = 2/√π·exp(-u²)` and `erf''(u) = -2u·erf'(u)`.
68    Erf(usize),
69    /// `a·ln(a)` with the `a = 0` limit `0` (GAMS `entropy` up to sign).
70    ///
71    /// Like [`TapeOp::Erf`] it has no AMPL `.nl` opcode and is reachable only
72    /// through the in-memory `Expr` path (`UnaryOp::XLogX`). Unlike `Erf` it is
73    /// fused for **correctness**, not convenience: the second derivative is
74    /// `1/a`, which is finite for every positive `a` down to `1e-308`, but
75    /// *every* chain-rule decomposition routes through `ln''(a) = -1/a²`. At
76    /// `a = 1e-299` that intermediate is `-1e598` — past `f64::MAX` — so the
77    /// composed Hessian is `inf`/`nan` while the answer it is computing,
78    /// `1e299`, is an ordinary number. No amount of care in the `Mul`/`Log`
79    /// rules reaches it; only an op that never forms `1/a²` does.
80    XLogX(usize),
81    /// `a·ln(a/b)` with the `a = 0` limit `0` — GAMS `centropy`. Operands
82    /// `(a, b)`; `b > 0` is the caller's responsibility, as for `Log`.
83    ///
84    /// Fused for [`TapeOp::XLogX`]'s reason (`∂²/∂a² = 1/a` is unreachable
85    /// through `ln''`) plus one of its own: `∂²/∂b² = a/b²` squares the
86    /// denominator, which overflows for `|b| > 1.3e154` even when `a/b²` is
87    /// itself in range. The fused rule computes `q = a/b` once and expresses
88    /// every second-order term as a division by `b`, never by `b²`.
89    CEntropy(usize, usize),
90    /// Two-argument arctangent `atan2(vals[a], vals[b])` (operands are
91    /// `(y, x)`, matching AMPL's `atan2(y, x)` / `.nl` opcode o48).
92    Atan2(usize, usize),
93    /// Pairwise minimum `min(vals[a], vals[b])`. Piecewise linear: the
94    /// value/tangent/adjoint route through whichever operand is smaller
95    /// (ties pick the first), and the second derivative is identically
96    /// zero. n-ary AMPL `min` (opcode o11) folds to a chain of these.
97    Min(usize, usize),
98    /// Pairwise maximum `max(vals[a], vals[b])` — the `Min` mirror;
99    /// n-ary AMPL `max` (opcode o12) folds to a chain of these.
100    Max(usize, usize),
101    /// Relational comparison `vals[a] OP vals[b]` → `1.0`/`0.0`.
102    /// Piecewise constant, so its derivative is identically zero — the
103    /// AD passes treat it as a constant w.r.t. its operands.
104    Cmp(CmpOp, usize, usize),
105    /// Logical AND: `1.0` iff both operands are nonzero. Zero derivative.
106    And(usize, usize),
107    /// Logical OR: `1.0` iff either operand is nonzero. Zero derivative.
108    Or(usize, usize),
109    /// Logical NOT: `1.0` iff the operand is zero. Zero derivative.
110    Not(usize),
111    /// `if-then-else`: operands `(cond, then, else)`. The value is
112    /// `vals[then]` when `vals[cond] != 0` else `vals[else]`, and the
113    /// value/tangent/adjoint all route through the active branch only.
114    /// The condition contributes no derivative (the branch switch is a
115    /// non-smooth event the AD ignores).
116    Select(usize, usize, usize),
117    /// AMPL imported (external) function call. The payload (library
118    /// handle, name, and argument list) is boxed so this rare variant
119    /// does not inflate `size_of::<TapeOp>()`: without the box the
120    /// `Arc`+`String`+`Vec` make every op ~64 bytes, which on a
121    /// summand-split objective with millions of tiny tapes (e.g.
122    /// `sensors`) costs gigabytes. Boxing drops the common arithmetic
123    /// ops back to the size of the next-largest variant.
124    Funcall(Box<FuncallData>),
125}
126
127/// Boxed payload of [`TapeOp::Funcall`]. The library is kept alive by
128/// the `Arc`; `name` is the registered function name; `args` carries
129/// positional arguments where real-valued args reference earlier tape
130/// slots and string args are inline literals.
131#[derive(Debug, Clone)]
132pub struct FuncallData {
133    pub lib: Arc<ExternalLibrary>,
134    pub name: String,
135    pub args: Vec<TapeFuncallArg>,
136}
137
138/// One argument of a `TapeOp::Funcall`. Real arguments are tape-slot indices
139/// (their values come from the running `vals[]` during forward); string
140/// arguments are owned literals (AMPL `h<len>:<chars>` tokens).
141#[derive(Debug, Clone)]
142pub enum TapeFuncallArg {
143    Tape(usize),
144    Str(String),
145}
146
147/// `ln(a/b)`, computed without ever materializing an out-of-range `a/b`.
148///
149/// Shared by every arm of [`TapeOp::CEntropy`] so the value and the derivatives
150/// cannot drift apart. Three regimes:
151///
152///   * `a ≈ b` — the ratio is near 1, where `ln` loses digits to cancellation.
153///     `(a - b)` is exact by Sterbenz whenever `b/2 ≤ a ≤ 2b`, so `ln_1p` of
154///     `(a - b)/b` keeps full precision where `q.ln()` would not.
155///   * `a/b` finite and positive — plain `q.ln()`.
156///   * `a/b` overflowed, underflowed to zero, or went non-finite — fall back to
157///     `ln(a) - ln(b)`. At `a = 1e300, b = 1e-300` the ratio is `1e600` (inf)
158///     but `ln` of it is a perfectly ordinary `1381.55`; the difference form
159///     reaches it. This is the only regime where cancellation is a risk, and
160///     it cannot bite here because the ratio being out of range means the two
161///     logs are far apart.
162#[inline]
163pub(crate) fn ln_ratio(a: f64, b: f64) -> f64 {
164    let q = a / b;
165    if q.is_finite() && q > 0.0 {
166        let t = (a - b) / b;
167        if t.abs() < 0.5 { t.ln_1p() } else { q.ln() }
168    } else {
169        a.ln() - b.ln()
170    }
171}
172
173/// `a·ln(a)`, the value arm of [`TapeOp::XLogX`] — GAMS `entropy`.
174///
175/// Exists as a FUSED op because the chain rule provably cannot evaluate its
176/// second derivative. `(a·ln a)'' = 1/a`, which at `a = 1e-299` is `1e299` and
177/// perfectly representable — but any decomposition routes through
178/// `ln''(a) = -1/a² = -1e598`, which exceeds `f64::MAX`. The composite is in
179/// range while the factor it is built from is not, so no rule for `Mul` and
180/// `Log` can recover it, however carefully written. Only a fused op that never
181/// forms `1/a²` gets there.
182///
183/// `0·ln 0` is `0` by the limit, not `NaN`.
184#[inline]
185pub(crate) fn xlogx(a: f64) -> f64 {
186    if a == 0.0 { 0.0 } else { a * a.ln() }
187}
188
189/// `a·ln(a/b)`, the value arm of [`TapeOp::CEntropy`] — GAMS `centropy`.
190///
191/// Fused for the same reason as [`xlogx`], plus one of its own: `∂²/∂b²` is
192/// `a/b²`, and `b²` overflows for `|b| > 1.3e154` while `a/b²` itself stays in
193/// range. The fused rule computes it as `q/b` and never squares anything.
194///
195/// `0·ln(0/b)` is `0` by the limit.
196#[inline]
197pub(crate) fn centropy(a: f64, b: f64) -> f64 {
198    if a == 0.0 { 0.0 } else { a * ln_ratio(a, b) }
199}
200
201/// `(a·ln a)' = ln(a) + 1`. `-inf` at `a = 0`, which is the true one-sided limit.
202#[inline]
203pub(crate) fn xlogx_d1(a: f64) -> f64 {
204    a.ln() + 1.0
205}
206
207/// `(a·ln a)'' = 1/a`.
208///
209/// The entire reason [`TapeOp::XLogX`] exists: this is finite for every positive
210/// `a`, while the `ln''` any decomposition would go through is not.
211#[inline]
212pub(crate) fn xlogx_d2(a: f64) -> f64 {
213    1.0 / a
214}
215
216/// `∂/∂a [a·ln(a/b)] = ln(a/b) + 1`.
217#[inline]
218pub(crate) fn centropy_da(a: f64, b: f64) -> f64 {
219    ln_ratio(a, b) + 1.0
220}
221
222/// `∂/∂b [a·ln(a/b)] = -a/b`.
223#[inline]
224pub(crate) fn centropy_db(a: f64, b: f64) -> f64 {
225    -(a / b)
226}
227
228/// `∂²/∂a² [a·ln(a/b)] = 1/a` — independent of `b`; see [`xlogx_d2`].
229#[inline]
230pub(crate) fn centropy_daa(a: f64) -> f64 {
231    1.0 / a
232}
233
234/// `∂²/∂a∂b [a·ln(a/b)] = -1/b` — independent of `a`.
235#[inline]
236pub(crate) fn centropy_dab(b: f64) -> f64 {
237    -1.0 / b
238}
239
240/// `∂²/∂b² [a·ln(a/b)] = a/b²`, evaluated as `(a/b)/b`.
241///
242/// Never as `a/(b*b)`: at `a = 1e300, b = 1e200` the true value is `1e-100`, but
243/// `b*b` is `inf` and the squared form returns `0.0`. Dividing twice keeps it.
244#[inline]
245pub(crate) fn centropy_dbb(a: f64, b: f64) -> f64 {
246    (a / b) / b
247}
248
249/// Inverse hyperbolic tangent, the value arm of [`TapeOp::Atanh`] (gh#1151).
250///
251/// Delegates to `libm` (fdlibm's algorithm) rather than `f64::atanh`. The
252/// standard library computes `0.5 * ln_1p(2u / (1 - u))`, and near `u = -1`
253/// the argument of `ln_1p` is `-1 + δ` with `δ` lost to rounding: relative
254/// error 3.6e-5 at `u = -1 + 1e-13` and 3.0e-3 at `u = -1 + 1e-15`, enough
255/// for a downstream global solver to certify the wrong objective.
256#[inline]
257pub(crate) fn atanh(u: f64) -> f64 {
258    libm::atanh(u)
259}
260
261// ---------------------------------------------------------------------------
262// Inverse-trig / hyperbolic derivative kernels (gh#1151 and its neighbours).
263//
264// Every AD sweep — reverse, forward-tangent, second-order, on both the `Tape`
265// and the `GlobalTape` paths, plus the `.nl` reader's derivative table — calls
266// these, so each formula is written (and tested) once. The textbook forms they
267// replace lose accuracy in one of two ways:
268//
269// * Cancellation. `1 - u*u` near `|u| = 1` and `u*u - 1` near `u = 1` subtract
270//   two nearly equal numbers after `u*u` has already rounded. For `u = 1 - δ`
271//   that rounding error is ~`δ²` until it reaches an ulp, so the relative
272//   error in `1 - u*u` peaks at ~4e-9 around `δ ≈ 3e-9` (measured against
273//   100-digit decimal) — small, but it is the gradient of a term a global
274//   solver is bounding. Factored as `(1-u)(1+u)` / `(u-1)(u+1)`, the factor
275//   that approaches zero is computed exactly (Sterbenz) and the product
276//   keeps full relative accuracy.
277//   `tanh'` written as `1 - tanh(u)²` cancels the same way once `tanh(u)` has
278//   rounded toward ±1 (rel. error ~1e-8 at `u = 10`, exactly 0 past `|u| ≈ 19`);
279//   `sech²(u)` does not.
280// * Overflow of an intermediate square. `u*u ± 1` is `inf` past
281//   `|u| ≈ 1.3e154`, so `1/sqrt(u*u ± 1)` returned `0` where the true value is
282//   `1/|u|`; `d*d` in the second derivatives overflows sooner still (`atan''`
283//   at `|u| > 1e77`). Second derivatives are therefore built from the first as
284//   `(u·g)·g·g` — `u·g` is O(1) — rather than as `u / (s·sqrt(s))`.
285// ---------------------------------------------------------------------------
286
287/// `atan'(u) = 1 / (1 + u²)`.
288///
289/// `1 + u*u` has no cancellation, but `u*u` overflows past `|u| ≈ 1.3e154`
290/// while the true value stays nonzero (subnormal) up to `|u| ≈ 6e161`. Past
291/// `|u| = 1e150` it is therefore taken as `r²` with `r = 1/u`: there
292/// `1 + r²` rounds to exactly `1`, so `r²` is the same quantity, and it
293/// underflows gradually instead of jumping to `0`. Below the threshold the
294/// textbook form is kept bit for bit.
295#[inline]
296pub(crate) fn atan_d1(u: f64) -> f64 {
297    if u.abs() > 1e150 {
298        let r = 1.0 / u;
299        r * r
300    } else {
301        1.0 / (1.0 + u * u)
302    }
303}
304
305/// `atan''(u) = -2u / (1 + u²)²`, as `-2·(u·g)·g`.
306#[inline]
307pub(crate) fn atan_d2(u: f64) -> f64 {
308    let g = atan_d1(u);
309    -2.0 * (u * g) * g
310}
311
312/// `asin'(u) = 1 / sqrt(1 - u²)`, with `1 - u²` as `(1-u)(1+u)`.
313#[inline]
314pub(crate) fn asin_d1(u: f64) -> f64 {
315    1.0 / ((1.0 - u) * (1.0 + u)).sqrt()
316}
317
318/// `asin''(u) = u / (1 - u²)^{3/2}`, as `(u·g)·g·g`.
319#[inline]
320pub(crate) fn asin_d2(u: f64) -> f64 {
321    let g = asin_d1(u);
322    (u * g) * g * g
323}
324
325/// `acos'(u) = -asin'(u)`.
326#[inline]
327pub(crate) fn acos_d1(u: f64) -> f64 {
328    -asin_d1(u)
329}
330
331/// `acos''(u) = -asin''(u)`.
332#[inline]
333pub(crate) fn acos_d2(u: f64) -> f64 {
334    -asin_d2(u)
335}
336
337/// `tanh'(u) = sech²(u)`, never `1 - tanh(u)²` (see the section comment).
338/// `1/cosh` underflows to `0` only where the true value does.
339#[inline]
340pub(crate) fn tanh_d1(u: f64) -> f64 {
341    let s = 1.0 / u.cosh();
342    s * s
343}
344
345/// `acosh'(u) = 1 / sqrt(u² - 1)`, as `1 / (sqrt(u-1)·sqrt(u+1))`: exact
346/// factor near `u = 1`, and no `u*u` to overflow at large `u`.
347#[inline]
348pub(crate) fn acosh_d1(u: f64) -> f64 {
349    1.0 / ((u - 1.0).sqrt() * (u + 1.0).sqrt())
350}
351
352/// `acosh''(u) = -u / (u² - 1)^{3/2}`, as `-(u·g)·g·g`.
353#[inline]
354pub(crate) fn acosh_d2(u: f64) -> f64 {
355    let g = acosh_d1(u);
356    -(u * g) * g * g
357}
358
359/// `asinh'(u) = 1 / sqrt(u² + 1)`, via `hypot` so `u*u` cannot overflow.
360#[inline]
361pub(crate) fn asinh_d1(u: f64) -> f64 {
362    1.0 / u.hypot(1.0)
363}
364
365/// `asinh''(u) = -u / (u² + 1)^{3/2}`, as `-(u·g)·g·g`.
366#[inline]
367pub(crate) fn asinh_d2(u: f64) -> f64 {
368    let g = asinh_d1(u);
369    -(u * g) * g * g
370}
371
372/// `atanh'(u) = 1 / (1 - u²)`, with `1 - u²` as `(1-u)(1+u)`.
373#[inline]
374pub(crate) fn atanh_d1(u: f64) -> f64 {
375    1.0 / ((1.0 - u) * (1.0 + u))
376}
377
378/// `atanh''(u) = 2u / (1 - u²)²`, as `2·(u·g)·g`.
379#[inline]
380pub(crate) fn atanh_d2(u: f64) -> f64 {
381    let g = atanh_d1(u);
382    2.0 * (u * g) * g
383}
384
385// The four kernels below are `pub` because `pounce-cli`'s `HessianProgram`
386// lowers the same ops and must not carry its own copies: it held the
387// pre-fix forms of all of them (and of division's) after the tape was fixed.
388// They are `#[doc(hidden)]`: shared with a sibling crate, not part of
389// `pounce-nl`'s documented interface.
390
391/// `log''(u) = -1/u²`. Past `|u| = 1e150` as `-(1/u)/u`: `u*u` overflows at
392/// `|u| ≈ 1.3e154` and read `0` where the true value is subnormal (nonzero to
393/// `|u| ≈ 1.3e162`). Below the threshold the old form is kept bit for bit.
394#[doc(hidden)]
395#[inline]
396pub fn log_d2(u: f64) -> f64 {
397    if u.abs() > 1e150 {
398        -(1.0 / u) / u
399    } else {
400        -1.0 / (u * u)
401    }
402}
403
404/// `x / (u·ln 10)`, i.e. `x · log10'(u)`. Past `|u| = 1e300` as
405/// `(x/u)/ln 10`: `u·ln 10` overflows at `|u| ≈ 7.8e307`. It takes the
406/// numerator so that below the threshold every arm still performs the single
407/// division it always did, bit for bit.
408#[doc(hidden)]
409#[inline]
410pub fn div_u_ln10(x: f64, u: f64) -> f64 {
411    let c = std::f64::consts::LN_10;
412    if u.abs() > 1e300 {
413        (x / u) / c
414    } else {
415        x / (u * c)
416    }
417}
418
419/// `log10''(u) = -1/(u²·ln 10)`, with the same threshold as [`log_d2`].
420#[doc(hidden)]
421#[inline]
422pub fn log10_d2(u: f64) -> f64 {
423    let c = std::f64::consts::LN_10;
424    if u.abs() > 1e150 {
425        -((1.0 / u) / u) / c
426    } else {
427        -1.0 / (u * u * c)
428    }
429}
430
431/// `sqrt''(u) = -1/(4·u^{3/2})`, given `sv = sqrt(u)`. Past `u = 1e200` as
432/// `(-0.25/u)/sv`: `u·sv` overflows at `u ≈ 2.6e205`, where the true value is
433/// still ~1e-308. Below the threshold the old form is kept bit for bit.
434#[doc(hidden)]
435#[inline]
436pub fn sqrt_d2(u: f64, sv: f64) -> f64 {
437    if u > 1e200 {
438        (-0.25 / u) / sv
439    } else {
440        -0.25 / (u * sv)
441    }
442}
443
444/// `(y/k, x/k, hypot(y/k, x/k), k)` for the `atan2` derivative kernels:
445/// `k = 2` once either argument passes `1e300`, else `k = 1`.
446///
447/// Near `f64::MAX` the plain forms overflow twice over: `hypot(y, x)` itself
448/// (`hypot(1.5e308, 1e308) ≈ 1.8e308`, so `h = inf` and a subnormal gradient
449/// reads `0`), and `y ± x` in the Hessian (`inf` beside `y = -x`, and
450/// `inf · 0 = NaN` on it). Halving a value that large is exact, so `y/k` and
451/// `x/k` carry the same ratios; the kernels divide by `k` again at the end.
452/// With `k = 1` every kernel is bit-identical to the unscaled form.
453#[inline]
454fn atan2_frame(y: f64, x: f64) -> (f64, f64, f64, f64) {
455    let k = if y.abs().max(x.abs()) > 1e300 {
456        2.0
457    } else {
458        1.0
459    };
460    let (ys, xs) = (y / k, x / k);
461    (ys, xs, ys.hypot(xs), k)
462}
463
464/// `∇atan2(y, x) = (x, -y) / (x² + y²)`, scaled by `h = hypot(y, x)`.
465///
466/// Never via `d = y*y + x*x`: `d` overflows past `|·| ≈ 1.3e154` (gradient
467/// `0` where the true one is `~1/h`) and underflows below `~1e-162` (`inf`).
468/// `(x/h)/h` keeps every intermediate O(1) or O(1/h). See [`atan2_frame`]
469/// for arguments near `f64::MAX`.
470#[inline]
471pub(crate) fn atan2_d1(y: f64, x: f64) -> (f64, f64) {
472    let (ys, xs, hs, k) = atan2_frame(y, x);
473    let (cy, cx) = (ys / hs, xs / hs);
474    // `/ k` while the quotient is still O(1), so it is exact; `/ hs` last,
475    // where a subnormal result rounds once.
476    ((cx / k) / hs, (-cy / k) / hs)
477}
478
479/// Hessian of `atan2(y, x)`: `(∂yy, ∂yx, ∂xx) = (-2xy, y² - x², 2xy) / (x² + y²)²`,
480/// with the same `hypot` scaling as [`atan2_d1`] and `y² - x²` factored as
481/// `(y-x)(y+x)` so it does not cancel near `y = ±x`.
482#[inline]
483pub(crate) fn atan2_d2(y: f64, x: f64) -> (f64, f64, f64) {
484    let (ys, xs, hs, k) = atan2_frame(y, x);
485    let (cy, cx) = (ys / hs, xs / hs);
486    let kk = k * k;
487    let q = 2.0 * cy * cx / kk / hs / hs;
488    // `y - x` before scaling, where it is exact near `y = x` (Sterbenz);
489    // `y/h - x/h` would subtract two already-rounded quotients.
490    let r = ((ys - xs) / hs) * ((ys + xs) / hs);
491    (-q, r / kk / hs / hs, q)
492}
493
494/// Gauss error function, the value arm of [`TapeOp::Erf`].
495///
496/// Delegates to `libm` (the rust-lang port of musl's libm) rather than a
497/// series approximation: `erf` shows up inside residuals, so a 1e-7-accurate
498/// Abramowitz–Stegun fit would cap the achievable KKT error well above what
499/// the IPM asks for.
500#[inline]
501pub(crate) fn erf(u: f64) -> f64 {
502    libm::erf(u)
503}
504
505/// `erf'(u) = 2/√π · exp(-u²)`.
506///
507/// The AD sweeps all call this (rather than each inlining the constant) so
508/// the tangent and adjoint arms of [`TapeOp::Erf`] cannot drift apart.
509#[inline]
510pub(crate) fn erf_d1(u: f64) -> f64 {
511    std::f64::consts::FRAC_2_SQRT_PI * (-u * u).exp()
512}
513
514/// `erf''(u) = -2u · erf'(u)`.
515///
516/// Parenthesized as `-2·(u·erf'(u))`, not `(-2·u)·erf'(u)`. The two differ
517/// at the top of the range: `erf_d1` underflows to `0.0` around `|u| > 27`,
518/// while `-2.0 * u` overflows to `±inf` for `|u| > f64::MAX/2`, and
519/// `inf * 0.0` is `NaN` where the true limit is `0`. Multiplying `u` into
520/// the already-underflowed derivative first keeps every finite magnitude
521/// finite. (`u = ±inf` is `NaN` either way, correctly — the model has
522/// already left the reals by then.) Only the Hessian arms read this; the
523/// value and gradient arms are unaffected.
524#[inline]
525pub(crate) fn erf_d2(u: f64) -> f64 {
526    -2.0 * (u * erf_d1(u))
527}
528
529/// Evaluate a relational opcode on two scalar values, returning the
530/// boolean truth (callers map it to `1.0`/`0.0`).
531#[inline]
532fn cmp_holds(op: CmpOp, a: f64, b: f64) -> bool {
533    match op {
534        CmpOp::Lt => a < b,
535        CmpOp::Le => a <= b,
536        CmpOp::Eq => a == b,
537        CmpOp::Ge => a >= b,
538        CmpOp::Gt => a > b,
539        CmpOp::Ne => a != b,
540    }
541}
542
543fn funcall_to_ext_args<'a>(args: &'a [TapeFuncallArg], vals: &[f64]) -> Vec<ExternalArg<'a>> {
544    args.iter()
545        .map(|a| match a {
546            TapeFuncallArg::Tape(idx) => ExternalArg::Real(vals[*idx]),
547            TapeFuncallArg::Str(s) => ExternalArg::Str(s.as_str()),
548        })
549        .collect()
550}
551
552/// Evaluate an external (AMPL imported) function, poisoning the result with
553/// `NaN` instead of panicking when the library reports an error.
554///
555/// An external eval fails on user-controllable conditions — most commonly an
556/// out-of-domain property evaluation (e.g. an IDAES Helmholtz thermo call
557/// outside its valid pressure/temperature range). We mirror the tape's own
558/// arithmetic domain-error semantics (`log(-1) → NaN`): hand back NaN so the
559/// IPM sees a failed evaluation and the line search backs off, rather than
560/// raising an uncatchable panic across the pyo3 boundary on the `read_nl`
561/// surface. The NaN derivative/Hessian vectors are sized by the full argument
562/// count — an upper bound on the real-arg count a successful eval returns — so
563/// every downstream index into them stays in range.
564fn ext_eval_or_nan(
565    lib: &ExternalLibrary,
566    name: &str,
567    call_args: &[ExternalArg<'_>],
568    n_args: usize,
569    want_derivs: bool,
570    want_hes: bool,
571) -> EvalResult {
572    lib.eval(name, call_args, want_derivs, want_hes)
573        .unwrap_or_else(|_| EvalResult {
574            value: f64::NAN,
575            derivs: want_derivs.then(|| vec![f64::NAN; n_args]),
576            hessian: want_hes.then(|| vec![f64::NAN; n_args * (n_args + 1) / 2]),
577        })
578}
579
580/// A flattened expression tape. The result of evaluation is the value
581/// at slot `ops.len() - 1` (i.e. the last op).
582#[derive(Debug, Clone)]
583pub struct Tape {
584    pub ops: Vec<TapeOp>,
585}
586
587impl Tape {
588    /// Build a tape from an `Expr` tree (no AMPL external functions). CSE
589    /// bodies (`Expr::Cse(rc)`) are cached by `Arc` pointer identity so each
590    /// body is emitted once even when referenced many times.
591    pub fn build(expr: &Expr) -> Self {
592        Self::build_with_externals(expr, &ExternalResolver::default())
593    }
594
595    /// Build a tape from an `Expr` tree, resolving any `Expr::Funcall`
596    /// nodes through `resolver`. Panics if the expression references a
597    /// funcall id that is not in the resolver — `NlProblem::resolve_externals`
598    /// must populate the resolver before tape construction.
599    pub fn build_with_externals(expr: &Expr, resolver: &ExternalResolver) -> Self {
600        let mut ops = Vec::new();
601        let mut cache: HashMap<*const Expr, usize> = HashMap::new();
602        build_recursive(expr, &mut ops, &mut cache, resolver);
603        Tape { ops }
604    }
605
606    /// Forward sweep: returns `vals[i] = value of tape slot i`. The
607    /// scalar tape result is `vals[ops.len() - 1]`.
608    pub fn forward(&self, x: &[f64]) -> Vec<f64> {
609        let mut vals: Vec<f64> = Vec::with_capacity(self.ops.len());
610        for op in &self.ops {
611            let v = match op {
612                TapeOp::Const(c) => *c,
613                TapeOp::Var(i) => x[*i],
614                TapeOp::Add(a, b) => vals[*a] + vals[*b],
615                TapeOp::Sub(a, b) => vals[*a] - vals[*b],
616                TapeOp::Mul(a, b) => vals[*a] * vals[*b],
617                TapeOp::Div(a, b) => vals[*a] / vals[*b],
618                TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
619                TapeOp::Neg(a) => -vals[*a],
620                TapeOp::Abs(a) => vals[*a].abs(),
621                TapeOp::Sqrt(a) => vals[*a].sqrt(),
622                TapeOp::Exp(a) => vals[*a].exp(),
623                TapeOp::Log(a) => vals[*a].ln(),
624                TapeOp::Log10(a) => vals[*a].log10(),
625                TapeOp::Sin(a) => vals[*a].sin(),
626                TapeOp::Cos(a) => vals[*a].cos(),
627                TapeOp::Tan(a) => vals[*a].tan(),
628                TapeOp::Atan(a) => vals[*a].atan(),
629                TapeOp::Acos(a) => vals[*a].acos(),
630                TapeOp::Sinh(a) => vals[*a].sinh(),
631                TapeOp::Cosh(a) => vals[*a].cosh(),
632                TapeOp::Tanh(a) => vals[*a].tanh(),
633                TapeOp::Asin(a) => vals[*a].asin(),
634                TapeOp::Acosh(a) => vals[*a].acosh(),
635                TapeOp::Asinh(a) => vals[*a].asinh(),
636                TapeOp::Atanh(a) => atanh(vals[*a]),
637                TapeOp::Erf(a) => erf(vals[*a]),
638                TapeOp::XLogX(a) => xlogx(vals[*a]),
639                TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
640                TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
641                TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
642                TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
643                TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
644                TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
645                TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
646                TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
647                TapeOp::Select(c, t, e) => {
648                    if vals[*c] != 0.0 {
649                        vals[*t]
650                    } else {
651                        vals[*e]
652                    }
653                }
654                TapeOp::Funcall(fc) => {
655                    let FuncallData { lib, name, args } = fc.as_ref();
656                    let call_args = funcall_to_ext_args(args, &vals);
657                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
658                    res.value
659                }
660            };
661            vals.push(v);
662        }
663        vals
664    }
665
666    pub fn eval(&self, x: &[f64]) -> f64 {
667        let vals = self.forward(x);
668        *vals.last().unwrap_or(&0.0)
669    }
670
671    /// Reverse-mode AD: accumulate `seed * df/dx_i` into `grad[i]` for
672    /// every problem variable `i` referenced by the tape. `grad` is
673    /// **not** zeroed by this routine — the caller can chain multiple
674    /// gradient accumulations into the same buffer.
675    pub fn gradient_seed(&self, x: &[f64], seed: f64, grad: &mut [f64]) {
676        if seed == 0.0 || self.ops.is_empty() {
677            return;
678        }
679        let vals = self.forward(x);
680        self.reverse(&vals, seed, grad);
681    }
682
683    /// Reverse-mode AD reusing two caller-supplied scratch buffers
684    /// (`vals` from [`forward_into`], and an `adj` arena ≥
685    /// `self.ops.len()`) instead of allocating a forward-value vector and
686    /// an adjoint vector per call like [`gradient_seed`]. The `.nl` design
687    /// emits one tiny tape per summand — ~10⁶ on large models — so a single
688    /// `eval_jac_g` / `eval_grad_f` drives this millions of times and the
689    /// per-call allocation dominated. `grad` is accumulated into (not
690    /// zeroed); `adj` may be passed dirty (it is zeroed at the touched
691    /// slots internally).
692    ///
693    /// [`forward_into`]: Tape::forward_into
694    pub fn gradient_seed_into(
695        &self,
696        x: &[f64],
697        seed: f64,
698        grad: &mut [f64],
699        vals: &mut [f64],
700        adj: &mut [f64],
701    ) {
702        if seed == 0.0 || self.ops.is_empty() {
703            return;
704        }
705        debug_assert!(vals.len() >= self.ops.len());
706        self.forward_into(x, vals);
707        self.reverse_into(vals, seed, grad, adj);
708    }
709
710    fn reverse(&self, vals: &[f64], seed: f64, grad: &mut [f64]) {
711        let n = self.ops.len();
712        let mut adj = vec![0.0f64; n];
713        self.reverse_into(vals, seed, grad, &mut adj);
714    }
715
716    /// Reverse adjoint sweep into a caller-supplied `adj` scratch buffer
717    /// (length ≥ `self.ops.len()`), the allocation-free core of [`reverse`].
718    /// `adj` is zeroed over `[0, n)` internally, so a dirty arena is fine;
719    /// `grad` is accumulated into (not zeroed).
720    fn reverse_into(&self, vals: &[f64], seed: f64, grad: &mut [f64], adj: &mut [f64]) {
721        let n = self.ops.len();
722        debug_assert!(adj.len() >= n);
723        adj[..n].fill(0.0);
724        adj[n - 1] = seed;
725
726        for i in (0..n).rev() {
727            let a = adj[i];
728            if a == 0.0 {
729                continue;
730            }
731            match &self.ops[i] {
732                TapeOp::Const(_) => {}
733                TapeOp::Var(j) => {
734                    grad[*j] += a;
735                }
736                TapeOp::Add(l, r) => {
737                    adj[*l] += a;
738                    adj[*r] += a;
739                }
740                TapeOp::Sub(l, r) => {
741                    adj[*l] += a;
742                    adj[*r] -= a;
743                }
744                TapeOp::Mul(l, r) => {
745                    adj[*l] += a * vals[*r];
746                    adj[*r] += a * vals[*l];
747                }
748                TapeOp::Div(l, r) => {
749                    // ∂(l/r)/∂r = -q/r, NOT -l/r²: r*r overflows for |r| > 1.3e154
750                    // and underflows for |r| < 1.5e-154, zeroing an adjoint that is
751                    // itself representable. `vals[i]` is q = l/r, already computed.
752                    let rv = vals[*r];
753                    adj[*l] += a / rv;
754                    adj[*r] -= a * vals[i] / rv;
755                }
756                TapeOp::Pow(l, r) => {
757                    let lv = vals[*l];
758                    let rv = vals[*r];
759                    if rv != 0.0 {
760                        adj[*l] += a * rv * lv.powf(rv - 1.0);
761                    }
762                    if lv > 0.0 {
763                        adj[*r] += a * vals[i] * lv.ln();
764                    }
765                }
766                TapeOp::Neg(j) => {
767                    adj[*j] -= a;
768                }
769                TapeOp::Abs(j) => {
770                    if vals[*j] >= 0.0 {
771                        adj[*j] += a;
772                    } else {
773                        adj[*j] -= a;
774                    }
775                }
776                TapeOp::Sqrt(j) => {
777                    let sv = vals[i];
778                    if sv > 0.0 {
779                        adj[*j] += a * 0.5 / sv;
780                    }
781                }
782                TapeOp::Exp(j) => {
783                    adj[*j] += a * vals[i];
784                }
785                TapeOp::Log(j) => {
786                    adj[*j] += a / vals[*j];
787                }
788                TapeOp::Log10(j) => {
789                    adj[*j] += div_u_ln10(a, vals[*j]);
790                }
791                TapeOp::Sin(j) => {
792                    adj[*j] += a * vals[*j].cos();
793                }
794                TapeOp::Cos(j) => {
795                    adj[*j] -= a * vals[*j].sin();
796                }
797                TapeOp::Tan(j) => {
798                    let t = vals[i];
799                    adj[*j] += a * (1.0 + t * t);
800                }
801                TapeOp::Atan(j) => {
802                    adj[*j] += a * atan_d1(vals[*j]);
803                }
804                TapeOp::Acos(j) => {
805                    adj[*j] += a * acos_d1(vals[*j]);
806                }
807                TapeOp::Sinh(j) => {
808                    adj[*j] += a * vals[*j].cosh();
809                }
810                TapeOp::Cosh(j) => {
811                    adj[*j] += a * vals[*j].sinh();
812                }
813                TapeOp::Tanh(j) => {
814                    adj[*j] += a * tanh_d1(vals[*j]);
815                }
816                TapeOp::Asin(j) => {
817                    adj[*j] += a * asin_d1(vals[*j]);
818                }
819                TapeOp::Acosh(j) => {
820                    adj[*j] += a * acosh_d1(vals[*j]);
821                }
822                TapeOp::Asinh(j) => {
823                    adj[*j] += a * asinh_d1(vals[*j]);
824                }
825                TapeOp::Atanh(j) => {
826                    adj[*j] += a * atanh_d1(vals[*j]);
827                }
828                TapeOp::Erf(j) => {
829                    adj[*j] += a * erf_d1(vals[*j]);
830                }
831                TapeOp::XLogX(j) => {
832                    adj[*j] += a * xlogx_d1(vals[*j]);
833                }
834                TapeOp::CEntropy(l, r) => {
835                    adj[*l] += a * centropy_da(vals[*l], vals[*r]);
836                    adj[*r] += a * centropy_db(vals[*l], vals[*r]);
837                }
838                TapeOp::Atan2(l, r) => {
839                    let y = vals[*l];
840                    let x = vals[*r];
841                    let (fa, fb) = atan2_d1(y, x);
842                    adj[*l] += a * fa;
843                    adj[*r] += a * fb;
844                }
845                // min/max are piecewise linear: the adjoint flows to the
846                // selected operand only (ties pick the first, a valid
847                // subgradient choice).
848                TapeOp::Min(l, r) => {
849                    if vals[*l] <= vals[*r] {
850                        adj[*l] += a;
851                    } else {
852                        adj[*r] += a;
853                    }
854                }
855                TapeOp::Max(l, r) => {
856                    if vals[*l] >= vals[*r] {
857                        adj[*l] += a;
858                    } else {
859                        adj[*r] += a;
860                    }
861                }
862                // Comparisons and logical connectives are piecewise
863                // constant: zero derivative, so no adjoint propagates.
864                TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
865                // if-then-else: the adjoint flows entirely into the
866                // active branch; the condition gets none.
867                TapeOp::Select(c, t, e) => {
868                    if vals[*c] != 0.0 {
869                        adj[*t] += a;
870                    } else {
871                        adj[*e] += a;
872                    }
873                }
874                TapeOp::Funcall(fc) => {
875                    let FuncallData { lib, name, args } = fc.as_ref();
876                    let call_args = funcall_to_ext_args(args, vals);
877                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
878                    let derivs = res.derivs.expect("want_derivs=true returns derivs");
879                    let mut k = 0usize;
880                    for arg in args {
881                        if let TapeFuncallArg::Tape(idx) = arg {
882                            adj[*idx] += a * derivs[k];
883                            k += 1;
884                        }
885                    }
886                }
887            }
888        }
889    }
890
891    /// Sorted distinct problem-variable indices that the tape depends on.
892    pub fn variables(&self) -> Vec<usize> {
893        let mut s: BTreeSet<usize> = BTreeSet::new();
894        for op in &self.ops {
895            if let TapeOp::Var(j) = op {
896                s.insert(*j);
897            }
898        }
899        s.into_iter().collect()
900    }
901
902    /// Forward tangent sweep: `dot[i] = d(slot_i) / dx_{seed_var}`.
903    /// Caller-supplied `dot` buffer is overwritten in full; no zeroing
904    /// needed beforehand because every slot is written before it is
905    /// read (the loop walks forward and only reads earlier slots).
906    fn forward_tangent(&self, vals: &[f64], seed_var: usize, dot: &mut [f64]) {
907        let n = self.ops.len();
908        debug_assert_eq!(dot.len(), n);
909        for i in 0..n {
910            dot[i] = match &self.ops[i] {
911                TapeOp::Const(_) => 0.0,
912                TapeOp::Var(k) => {
913                    if *k == seed_var {
914                        1.0
915                    } else {
916                        0.0
917                    }
918                }
919                TapeOp::Add(a, b) => dot[*a] + dot[*b],
920                TapeOp::Sub(a, b) => dot[*a] - dot[*b],
921                TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
922                TapeOp::Div(a, b) => {
923                    // (ȧ - q·ḃ)/b, NOT (ȧ·b - a·ḃ)/b². Squaring the denominator
924                    // overflows for |b| > 1.3e154 and underflows for |b| < 1.5e-154,
925                    // destroying a tangent that is itself perfectly representable:
926                    // at a = b = 1e300 with ḃ = 1 the true -a/b² is -1e-300, but
927                    // b*b is inf so the squared form returns -0.0. Dividing by b
928                    // twice instead of by b² once is the standard Kahan form, and
929                    // costs nothing -- q is this node's own value, already computed.
930                    (dot[*a] - vals[i] * dot[*b]) / vals[*b]
931                }
932                TapeOp::Pow(a, b) => {
933                    let u = vals[*a];
934                    let r = vals[*b];
935                    let du = dot[*a];
936                    let dr = dot[*b];
937                    let mut result = 0.0;
938                    // Match the reverse-mode gradient's guard (`rv != 0.0` only): at base
939                    // u == 0 the slope is still well defined for r >= 1 (and a
940                    // genuine ±inf for r < 1), so it must not be silently dropped,
941                    // or the forward tangent disagrees with the reverse gradient.
942                    if r != 0.0 {
943                        result += r * u.powf(r - 1.0) * du;
944                    }
945                    if u > 0.0 {
946                        result += vals[i] * u.ln() * dr;
947                    }
948                    result
949                }
950                TapeOp::Neg(a) => -dot[*a],
951                TapeOp::Abs(a) => {
952                    if vals[*a] >= 0.0 {
953                        dot[*a]
954                    } else {
955                        -dot[*a]
956                    }
957                }
958                TapeOp::Sqrt(a) => {
959                    let sv = vals[i];
960                    if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
961                }
962                TapeOp::Exp(a) => dot[*a] * vals[i],
963                TapeOp::Log(a) => dot[*a] / vals[*a],
964                TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
965                TapeOp::Sin(a) => dot[*a] * vals[*a].cos(),
966                TapeOp::Cos(a) => -dot[*a] * vals[*a].sin(),
967                TapeOp::Tan(a) => {
968                    let t = vals[i];
969                    dot[*a] * (1.0 + t * t)
970                }
971                TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
972                TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
973                TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
974                TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
975                TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
976                TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
977                TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
978                TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
979                TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
980                TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
981                TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
982                TapeOp::CEntropy(a, b) => {
983                    centropy_da(vals[*a], vals[*b]) * dot[*a]
984                        + centropy_db(vals[*a], vals[*b]) * dot[*b]
985                }
986                TapeOp::Atan2(a, b) => {
987                    let y = vals[*a];
988                    let x = vals[*b];
989                    let (fa, fb) = atan2_d1(y, x);
990                    fa * dot[*a] + fb * dot[*b]
991                }
992                // min/max: the tangent follows the selected operand.
993                TapeOp::Min(a, b) => {
994                    if vals[*a] <= vals[*b] {
995                        dot[*a]
996                    } else {
997                        dot[*b]
998                    }
999                }
1000                TapeOp::Max(a, b) => {
1001                    if vals[*a] >= vals[*b] {
1002                        dot[*a]
1003                    } else {
1004                        dot[*b]
1005                    }
1006                }
1007                TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
1008                TapeOp::Select(c, t, e) => {
1009                    if vals[*c] != 0.0 {
1010                        dot[*t]
1011                    } else {
1012                        dot[*e]
1013                    }
1014                }
1015                TapeOp::Funcall(fc) => {
1016                    let FuncallData { lib, name, args } = fc.as_ref();
1017                    let call_args = funcall_to_ext_args(args, vals);
1018                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
1019                    let derivs = res.derivs.expect("want_derivs=true returns derivs");
1020                    let mut acc = 0.0;
1021                    let mut k = 0usize;
1022                    for arg in args {
1023                        if let TapeFuncallArg::Tape(idx) = arg {
1024                            acc += derivs[k] * dot[*idx];
1025                            k += 1;
1026                        }
1027                    }
1028                    acc
1029                }
1030            };
1031        }
1032    }
1033
1034    /// Forward sweep into a caller-supplied buffer. Avoids the
1035    /// per-call allocation of `forward()` so hot paths can reuse
1036    /// one scratch arena across many tapes.
1037    pub fn forward_into(&self, x: &[f64], vals: &mut [f64]) {
1038        let n = self.ops.len();
1039        debug_assert!(vals.len() >= n);
1040        for i in 0..n {
1041            vals[i] = match &self.ops[i] {
1042                TapeOp::Const(c) => *c,
1043                TapeOp::Var(j) => x[*j],
1044                TapeOp::Add(a, b) => vals[*a] + vals[*b],
1045                TapeOp::Sub(a, b) => vals[*a] - vals[*b],
1046                TapeOp::Mul(a, b) => vals[*a] * vals[*b],
1047                TapeOp::Div(a, b) => vals[*a] / vals[*b],
1048                TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
1049                TapeOp::Neg(a) => -vals[*a],
1050                TapeOp::Abs(a) => vals[*a].abs(),
1051                TapeOp::Sqrt(a) => vals[*a].sqrt(),
1052                TapeOp::Exp(a) => vals[*a].exp(),
1053                TapeOp::Log(a) => vals[*a].ln(),
1054                TapeOp::Log10(a) => vals[*a].log10(),
1055                TapeOp::Sin(a) => vals[*a].sin(),
1056                TapeOp::Cos(a) => vals[*a].cos(),
1057                TapeOp::Tan(a) => vals[*a].tan(),
1058                TapeOp::Atan(a) => vals[*a].atan(),
1059                TapeOp::Acos(a) => vals[*a].acos(),
1060                TapeOp::Sinh(a) => vals[*a].sinh(),
1061                TapeOp::Cosh(a) => vals[*a].cosh(),
1062                TapeOp::Tanh(a) => vals[*a].tanh(),
1063                TapeOp::Asin(a) => vals[*a].asin(),
1064                TapeOp::Acosh(a) => vals[*a].acosh(),
1065                TapeOp::Asinh(a) => vals[*a].asinh(),
1066                TapeOp::Atanh(a) => atanh(vals[*a]),
1067                TapeOp::Erf(a) => erf(vals[*a]),
1068                TapeOp::XLogX(a) => xlogx(vals[*a]),
1069                TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
1070                TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
1071                TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
1072                TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
1073                TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
1074                TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
1075                TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
1076                TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
1077                TapeOp::Select(c, t, e) => {
1078                    if vals[*c] != 0.0 {
1079                        vals[*t]
1080                    } else {
1081                        vals[*e]
1082                    }
1083                }
1084                TapeOp::Funcall(fc) => {
1085                    let FuncallData { lib, name, args } = fc.as_ref();
1086                    let call_args = funcall_to_ext_args(args, &*vals);
1087                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
1088                    res.value
1089                }
1090            };
1091        }
1092    }
1093
1094    /// Scalar tape value, reusing a caller-supplied scratch buffer
1095    /// (`vals.len() >= self.ops.len()`) instead of allocating a fresh
1096    /// forward-value `Vec` per call like [`eval`]. The `.nl` design emits
1097    /// one tiny tape per summand — 10⁵–10⁶ on large models — so a single
1098    /// `eval_f` / `eval_g` drives this once per summand and the per-call
1099    /// allocation dominated the sweep itself (same motivation as
1100    /// [`gradient_seed_into`], M18).
1101    ///
1102    /// [`eval`]: Tape::eval
1103    /// [`gradient_seed_into`]: Tape::gradient_seed_into
1104    pub fn eval_into(&self, x: &[f64], vals: &mut [f64]) -> f64 {
1105        let n = self.ops.len();
1106        if n == 0 {
1107            return 0.0;
1108        }
1109        self.forward_into(x, vals);
1110        vals[n - 1]
1111    }
1112
1113    /// Directional Hessian-vector product: emits
1114    /// `weight * (∇²f · seed)[k]` into `out[k]` for every problem
1115    /// variable `k` the tape references. Caller supplies the
1116    /// forward-pass result `vals` (use [`forward_into`]) plus three
1117    /// scratch buffers (`dot`, `adj`, `adj_dot`), each at least
1118    /// `self.ops.len()` long. `out` must be at least one past the
1119    /// largest variable index in the tape; the routine reads
1120    /// `seed[k]` for each `Var(k)` and writes `out[k] += weight *
1121    /// (Hess · seed)[k]`.
1122    ///
1123    /// This is one forward-over-reverse AD pass — O(n_ops) work —
1124    /// regardless of how many variables the tape depends on, which
1125    /// is what makes Hessian coloring efficient: a single
1126    /// directional pass recovers a whole color group of columns.
1127    ///
1128    /// [`forward_into`]: Tape::forward_into
1129    pub fn hessian_directional(
1130        &self,
1131        vals: &[f64],
1132        seed: &[f64],
1133        weight: f64,
1134        out: &mut [f64],
1135        dot: &mut [f64],
1136        adj: &mut [f64],
1137        adj_dot: &mut [f64],
1138    ) {
1139        let n = self.ops.len();
1140        if n == 0 || weight == 0.0 {
1141            return;
1142        }
1143        debug_assert!(vals.len() >= n);
1144        debug_assert!(dot.len() >= n);
1145        debug_assert!(adj.len() >= n);
1146        debug_assert!(adj_dot.len() >= n);
1147
1148        // Forward tangent: dot[i] = (∂vals[i] / ∂x · seed). At
1149        // Var(k) the seed entry feeds in; the rest of the chain
1150        // rule matches `forward_tangent` exactly.
1151        for i in 0..n {
1152            dot[i] = match &self.ops[i] {
1153                TapeOp::Const(_) => 0.0,
1154                TapeOp::Var(k) => seed[*k],
1155                TapeOp::Add(a, b) => dot[*a] + dot[*b],
1156                TapeOp::Sub(a, b) => dot[*a] - dot[*b],
1157                TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
1158                TapeOp::Div(a, b) => {
1159                    // (ȧ - q·ḃ)/b, NOT (ȧ·b - a·ḃ)/b². Squaring the denominator
1160                    // overflows for |b| > 1.3e154 and underflows for |b| < 1.5e-154,
1161                    // destroying a tangent that is itself perfectly representable:
1162                    // at a = b = 1e300 with ḃ = 1 the true -a/b² is -1e-300, but
1163                    // b*b is inf so the squared form returns -0.0. Dividing by b
1164                    // twice instead of by b² once is the standard Kahan form, and
1165                    // costs nothing -- q is this node's own value, already computed.
1166                    (dot[*a] - vals[i] * dot[*b]) / vals[*b]
1167                }
1168                TapeOp::Pow(a, b) => {
1169                    let u = vals[*a];
1170                    let r = vals[*b];
1171                    let du = dot[*a];
1172                    let dr = dot[*b];
1173                    let mut result = 0.0;
1174                    // Match the reverse-mode gradient's guard (`rv != 0.0` only): at base
1175                    // u == 0 the slope is still well defined for r >= 1 (and a
1176                    // genuine ±inf for r < 1), so it must not be silently dropped,
1177                    // or the forward tangent disagrees with the reverse gradient.
1178                    if r != 0.0 {
1179                        result += r * u.powf(r - 1.0) * du;
1180                    }
1181                    if u > 0.0 {
1182                        result += vals[i] * u.ln() * dr;
1183                    }
1184                    result
1185                }
1186                TapeOp::Neg(a) => -dot[*a],
1187                TapeOp::Abs(a) => {
1188                    if vals[*a] >= 0.0 {
1189                        dot[*a]
1190                    } else {
1191                        -dot[*a]
1192                    }
1193                }
1194                TapeOp::Sqrt(a) => {
1195                    let sv = vals[i];
1196                    if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
1197                }
1198                TapeOp::Exp(a) => vals[i] * dot[*a],
1199                TapeOp::Log(a) => dot[*a] / vals[*a],
1200                TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
1201                TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
1202                TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
1203                TapeOp::Tan(a) => {
1204                    let t = vals[i];
1205                    (1.0 + t * t) * dot[*a]
1206                }
1207                TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
1208                TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
1209                TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
1210                TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
1211                TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
1212                TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
1213                TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
1214                TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
1215                TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
1216                TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
1217                TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
1218                TapeOp::CEntropy(a, b) => {
1219                    centropy_da(vals[*a], vals[*b]) * dot[*a]
1220                        + centropy_db(vals[*a], vals[*b]) * dot[*b]
1221                }
1222                TapeOp::Atan2(a, b) => {
1223                    let y = vals[*a];
1224                    let x = vals[*b];
1225                    let (fa, fb) = atan2_d1(y, x);
1226                    fa * dot[*a] + fb * dot[*b]
1227                }
1228                // min/max: the tangent follows the selected operand.
1229                TapeOp::Min(a, b) => {
1230                    if vals[*a] <= vals[*b] {
1231                        dot[*a]
1232                    } else {
1233                        dot[*b]
1234                    }
1235                }
1236                TapeOp::Max(a, b) => {
1237                    if vals[*a] >= vals[*b] {
1238                        dot[*a]
1239                    } else {
1240                        dot[*b]
1241                    }
1242                }
1243                TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
1244                TapeOp::Select(c, t, e) => {
1245                    if vals[*c] != 0.0 {
1246                        dot[*t]
1247                    } else {
1248                        dot[*e]
1249                    }
1250                }
1251                TapeOp::Funcall(fc) => {
1252                    let FuncallData { lib, name, args } = fc.as_ref();
1253                    let call_args = funcall_to_ext_args(args, vals);
1254                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
1255                    let derivs = res.derivs.expect("want_derivs=true returns derivs");
1256                    let mut acc = 0.0;
1257                    let mut k = 0usize;
1258                    for arg in args {
1259                        if let TapeFuncallArg::Tape(idx) = arg {
1260                            acc += derivs[k] * dot[*idx];
1261                            k += 1;
1262                        }
1263                    }
1264                    acc
1265                }
1266            };
1267        }
1268
1269        // Reverse over tangent. adj[i] = ∂f/∂vals[i],
1270        // adj_dot[i] = derivative of adj[i] along `seed`
1271        // direction = (Hess · seed) projected onto slot i.
1272        for slot in adj.iter_mut().take(n) {
1273            *slot = 0.0;
1274        }
1275        for slot in adj_dot.iter_mut().take(n) {
1276            *slot = 0.0;
1277        }
1278        adj[n - 1] = 1.0;
1279
1280        for i in (0..n).rev() {
1281            let w = adj[i];
1282            let wd = adj_dot[i];
1283            if w == 0.0 && wd == 0.0 {
1284                continue;
1285            }
1286            match &self.ops[i] {
1287                TapeOp::Const(_) => {}
1288                TapeOp::Var(k) => {
1289                    if wd != 0.0 {
1290                        out[*k] += weight * wd;
1291                    }
1292                }
1293                TapeOp::Add(a, b) => {
1294                    adj[*a] += w;
1295                    adj[*b] += w;
1296                    adj_dot[*a] += wd;
1297                    adj_dot[*b] += wd;
1298                }
1299                TapeOp::Sub(a, b) => {
1300                    adj[*a] += w;
1301                    adj[*b] -= w;
1302                    adj_dot[*a] += wd;
1303                    adj_dot[*b] -= wd;
1304                }
1305                TapeOp::Mul(a, b) => {
1306                    adj[*a] += w * vals[*b];
1307                    adj[*b] += w * vals[*a];
1308                    adj_dot[*a] += wd * vals[*b] + w * dot[*b];
1309                    adj_dot[*b] += wd * vals[*a] + w * dot[*a];
1310                }
1311                TapeOp::Div(a, b) => {
1312                    // Kahan form, second order. `vb2`/`vb3` overflow for
1313                    // |b| > 1.3e154 (and vb2 underflows below 1.5e-154), which
1314                    // silently zeroed second-order terms that are representable.
1315                    // Every 1/b² here is regrouped as (·/b)/b, and the
1316                    // 2a·ḃ/b³ term folds into q̇ = dot[i], already computed:
1317                    //   -ȧ/b² + 2q·ḃ/b² == (-q̇ + q·(ḃ/b)) / b.
1318                    let vb = vals[*b];
1319                    let q = vals[i];
1320                    let qd = dot[i];
1321                    adj[*a] += w / vb;
1322                    adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1323                    adj[*b] -= w * q / vb;
1324                    adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1325                }
1326                TapeOp::Pow(a, b) => {
1327                    let u = vals[*a];
1328                    let r = vals[*b];
1329                    let du = dot[*a];
1330                    let dr = dot[*b];
1331                    if r != 0.0 {
1332                        if u != 0.0 {
1333                            let p_a = r * u.powf(r - 1.0);
1334                            adj[*a] += w * p_a;
1335                            let mut dp_a = dr * u.powf(r - 1.0);
1336                            if u > 0.0 {
1337                                dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1338                            } else {
1339                                dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1340                            }
1341                            adj_dot[*a] += wd * p_a + w * dp_a;
1342                        } else if r >= 2.0 {
1343                            let p_a = 0.0;
1344                            adj[*a] += w * p_a;
1345                            let dp_a = if r == 2.0 {
1346                                2.0 * du
1347                            } else {
1348                                r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1349                            };
1350                            adj_dot[*a] += wd * p_a + w * dp_a;
1351                        }
1352                    }
1353                    if u > 0.0 {
1354                        let ln_u = u.ln();
1355                        let p_b = vals[i] * ln_u;
1356                        adj[*b] += w * p_b;
1357                        let dur = vals[i] * (r * du / u + dr * ln_u);
1358                        let dp_b = dur * ln_u + vals[i] * du / u;
1359                        adj_dot[*b] += wd * p_b + w * dp_b;
1360                    }
1361                }
1362                TapeOp::Neg(a) => {
1363                    adj[*a] -= w;
1364                    adj_dot[*a] -= wd;
1365                }
1366                TapeOp::Abs(a) => {
1367                    let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
1368                    adj[*a] += w * s;
1369                    adj_dot[*a] += wd * s;
1370                }
1371                TapeOp::Sqrt(a) => {
1372                    let sv = vals[i];
1373                    if sv > 0.0 {
1374                        let fp = 0.5 / sv;
1375                        let fpp = sqrt_d2(vals[*a], sv);
1376                        adj[*a] += w * fp;
1377                        adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1378                    }
1379                }
1380                TapeOp::Exp(a) => {
1381                    let ev = vals[i];
1382                    adj[*a] += w * ev;
1383                    adj_dot[*a] += wd * ev + w * ev * dot[*a];
1384                }
1385                TapeOp::Log(a) => {
1386                    let u = vals[*a];
1387                    adj[*a] += w / u;
1388                    adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
1389                }
1390                TapeOp::Log10(a) => {
1391                    let u = vals[*a];
1392                    adj[*a] += div_u_ln10(w, u);
1393                    adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
1394                }
1395                TapeOp::Sin(a) => {
1396                    let u = vals[*a];
1397                    let cu = u.cos();
1398                    adj[*a] += w * cu;
1399                    adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1400                }
1401                TapeOp::Cos(a) => {
1402                    let u = vals[*a];
1403                    let su = u.sin();
1404                    adj[*a] -= w * su;
1405                    adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1406                }
1407                TapeOp::Tan(a) => {
1408                    let t = vals[i];
1409                    let gp = 1.0 + t * t;
1410                    let gpp = 2.0 * t * gp;
1411                    adj[*a] += w * gp;
1412                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1413                }
1414                TapeOp::Atan(a) => {
1415                    let u = vals[*a];
1416                    let gp = atan_d1(u);
1417                    let gpp = atan_d2(u);
1418                    adj[*a] += w * gp;
1419                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1420                }
1421                TapeOp::Acos(a) => {
1422                    let u = vals[*a];
1423                    let gp = acos_d1(u);
1424                    let gpp = acos_d2(u);
1425                    adj[*a] += w * gp;
1426                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1427                }
1428                TapeOp::Sinh(a) => {
1429                    let u = vals[*a];
1430                    let gp = u.cosh();
1431                    let gpp = u.sinh();
1432                    adj[*a] += w * gp;
1433                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1434                }
1435                TapeOp::Cosh(a) => {
1436                    let u = vals[*a];
1437                    let gp = u.sinh();
1438                    let gpp = u.cosh();
1439                    adj[*a] += w * gp;
1440                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1441                }
1442                TapeOp::Tanh(a) => {
1443                    let t = vals[i];
1444                    let gp = tanh_d1(vals[*a]);
1445                    let gpp = -2.0 * t * gp;
1446                    adj[*a] += w * gp;
1447                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1448                }
1449                TapeOp::Asin(a) => {
1450                    let u = vals[*a];
1451                    let gp = asin_d1(u);
1452                    let gpp = asin_d2(u);
1453                    adj[*a] += w * gp;
1454                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1455                }
1456                TapeOp::Acosh(a) => {
1457                    let u = vals[*a];
1458                    let gp = acosh_d1(u);
1459                    let gpp = acosh_d2(u);
1460                    adj[*a] += w * gp;
1461                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1462                }
1463                TapeOp::Asinh(a) => {
1464                    let u = vals[*a];
1465                    let gp = asinh_d1(u);
1466                    let gpp = asinh_d2(u);
1467                    adj[*a] += w * gp;
1468                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1469                }
1470                TapeOp::Atanh(a) => {
1471                    let u = vals[*a];
1472                    let gp = atanh_d1(u);
1473                    let gpp = atanh_d2(u);
1474                    adj[*a] += w * gp;
1475                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1476                }
1477                TapeOp::Erf(a) => {
1478                    let u = vals[*a];
1479                    let gp = erf_d1(u);
1480                    let gpp = erf_d2(u);
1481                    adj[*a] += w * gp;
1482                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1483                }
1484                TapeOp::XLogX(a) => {
1485                    // gpp is 1/u, NOT the -1/u² any decomposition through
1486                    // `ln''` would produce. That is the whole point of the op.
1487                    let u = vals[*a];
1488                    let gp = xlogx_d1(u);
1489                    let gpp = xlogx_d2(u);
1490                    adj[*a] += w * gp;
1491                    adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1492                }
1493                TapeOp::CEntropy(a, b) => {
1494                    let ua = vals[*a];
1495                    let ub = vals[*b];
1496                    let fa = centropy_da(ua, ub);
1497                    let fb = centropy_db(ua, ub);
1498                    let faa = centropy_daa(ua);
1499                    let fab = centropy_dab(ub);
1500                    let fbb = centropy_dbb(ua, ub);
1501                    adj[*a] += w * fa;
1502                    adj[*b] += w * fb;
1503                    adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1504                    adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1505                }
1506                TapeOp::Atan2(a, b) => {
1507                    let y = vals[*a];
1508                    let x = vals[*b];
1509                    let (fa, fb) = atan2_d1(y, x);
1510                    let (faa, fab, fbb) = atan2_d2(y, x);
1511                    adj[*a] += w * fa;
1512                    adj[*b] += w * fb;
1513                    adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1514                    adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1515                }
1516                // min/max are piecewise linear (zero second derivative):
1517                // route the adjoint and its tangent into the selected
1518                // operand, exactly like the active branch of a Select.
1519                TapeOp::Min(a, b) => {
1520                    let br = if vals[*a] <= vals[*b] { *a } else { *b };
1521                    adj[br] += w;
1522                    adj_dot[br] += wd;
1523                }
1524                TapeOp::Max(a, b) => {
1525                    let br = if vals[*a] >= vals[*b] { *a } else { *b };
1526                    adj[br] += w;
1527                    adj_dot[br] += wd;
1528                }
1529                // Zero derivative: no first- or second-order adjoint.
1530                TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
1531                // Route both the adjoint and its tangent into the
1532                // active branch; the condition contributes nothing.
1533                TapeOp::Select(c, t, e) => {
1534                    let br = if vals[*c] != 0.0 { *t } else { *e };
1535                    adj[br] += w;
1536                    adj_dot[br] += wd;
1537                }
1538                TapeOp::Funcall(fc) => {
1539                    let FuncallData { lib, name, args } = fc.as_ref();
1540                    let call_args = funcall_to_ext_args(args, vals);
1541                    let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1542                    let derivs = res.derivs.expect("want_derivs=true returns derivs");
1543                    let hes = res.hessian.expect("want_hes=true returns hessian");
1544                    let real_tape: Vec<usize> = args
1545                        .iter()
1546                        .filter_map(|a| match a {
1547                            TapeFuncallArg::Tape(t) => Some(*t),
1548                            TapeFuncallArg::Str(_) => None,
1549                        })
1550                        .collect();
1551                    for (k, &tk) in real_tape.iter().enumerate() {
1552                        adj[tk] += w * derivs[k];
1553                        let mut second_term = 0.0;
1554                        for (l, &tl) in real_tape.iter().enumerate() {
1555                            let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1556                            let h_kl = hes[lo + hi * (hi + 1) / 2];
1557                            second_term += h_kl * dot[tl];
1558                        }
1559                        adj_dot[tk] += wd * derivs[k] + w * second_term;
1560                    }
1561                }
1562            }
1563        }
1564    }
1565
1566    /// Forward-over-reverse Hessian: for each variable `j` the tape
1567    /// depends on, accumulate `weight * (d²f / dx_i dx_j)` into
1568    /// `values[hess_map[(i, j)]]` for every `(i, j)` lower-triangle
1569    /// pair in the map. The same routine is used for the objective
1570    /// (with `weight = obj_factor`) and each active constraint (with
1571    /// `weight = lambda[k]`); contributions sum into the shared map.
1572    pub fn hessian_accumulate(
1573        &self,
1574        x: &[f64],
1575        weight: f64,
1576        hess_map: &HashMap<(usize, usize), usize>,
1577        values: &mut [f64],
1578    ) {
1579        let n = self.ops.len();
1580        if n == 0 || weight == 0.0 {
1581            return;
1582        }
1583        let v = self.forward(x);
1584        let var_indices = self.variables();
1585
1586        // Hoist scratch allocations out of the per-variable loop —
1587        // each was costing O(n) per j on every hessian_accumulate
1588        // call, which dominated runtime on large tapes (the dense-
1589        // Hessian Mittelmann problems). `forward_tangent` fully
1590        // overwrites `dot`, so no reset is needed there. `adj` and
1591        // `adj_dot` are mutated additively, so we zero them per j.
1592        let mut dot = vec![0.0f64; n];
1593        let mut adj = vec![0.0f64; n];
1594        let mut adj_dot = vec![0.0f64; n];
1595        for &j in &var_indices {
1596            self.forward_tangent(&v, j, &mut dot);
1597
1598            // adj[i] = standard adjoint (∂f/∂slot_i)
1599            // adj_dot[i] = derivative of adj[i] w.r.t. x_j = ∂²f/(∂slot_i ∂x_j)
1600            adj.fill(0.0);
1601            adj_dot.fill(0.0);
1602            adj[n - 1] = 1.0;
1603
1604            for i in (0..n).rev() {
1605                let w = adj[i];
1606                let wd = adj_dot[i];
1607                if w == 0.0 && wd == 0.0 {
1608                    continue;
1609                }
1610                match &self.ops[i] {
1611                    TapeOp::Const(_) => {}
1612                    TapeOp::Var(k) => {
1613                        // Lower-triangle: only emit when row k >= col j
1614                        // so an off-diagonal pair appears once.
1615                        if wd != 0.0 && *k >= j {
1616                            if let Some(&pos) = hess_map.get(&(*k, j)) {
1617                                values[pos] += weight * wd;
1618                            }
1619                        }
1620                    }
1621                    TapeOp::Add(a, b) => {
1622                        adj[*a] += w;
1623                        adj[*b] += w;
1624                        adj_dot[*a] += wd;
1625                        adj_dot[*b] += wd;
1626                    }
1627                    TapeOp::Sub(a, b) => {
1628                        adj[*a] += w;
1629                        adj[*b] -= w;
1630                        adj_dot[*a] += wd;
1631                        adj_dot[*b] -= wd;
1632                    }
1633                    TapeOp::Mul(a, b) => {
1634                        adj[*a] += w * v[*b];
1635                        adj[*b] += w * v[*a];
1636                        adj_dot[*a] += wd * v[*b] + w * dot[*b];
1637                        adj_dot[*b] += wd * v[*a] + w * dot[*a];
1638                    }
1639                    TapeOp::Div(a, b) => {
1640                        // Kahan form -- see the identical arm in the dense
1641                        // reverse-over-forward sweep for the derivation.
1642                        let vb = v[*b];
1643                        let q = v[i];
1644                        let qd = dot[i];
1645                        adj[*a] += w / vb;
1646                        adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1647                        adj[*b] -= w * q / vb;
1648                        adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1649                    }
1650                    TapeOp::Pow(a, b) => {
1651                        let u = v[*a];
1652                        let r = v[*b];
1653                        let du = dot[*a];
1654                        let dr = dot[*b];
1655                        if r != 0.0 {
1656                            if u != 0.0 {
1657                                let p_a = r * u.powf(r - 1.0);
1658                                adj[*a] += w * p_a;
1659                                let mut dp_a = dr * u.powf(r - 1.0);
1660                                if u > 0.0 {
1661                                    dp_a +=
1662                                        r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1663                                } else {
1664                                    dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1665                                }
1666                                adj_dot[*a] += wd * p_a + w * dp_a;
1667                            } else if r >= 2.0 {
1668                                let p_a = 0.0;
1669                                adj[*a] += w * p_a;
1670                                let dp_a = if r == 2.0 {
1671                                    2.0 * du
1672                                } else {
1673                                    r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1674                                };
1675                                adj_dot[*a] += wd * p_a + w * dp_a;
1676                            }
1677                        }
1678                        if u > 0.0 {
1679                            let ln_u = u.ln();
1680                            let p_b = v[i] * ln_u;
1681                            adj[*b] += w * p_b;
1682                            let dur = v[i] * (r * du / u + dr * ln_u);
1683                            let dp_b = dur * ln_u + v[i] * du / u;
1684                            adj_dot[*b] += wd * p_b + w * dp_b;
1685                        }
1686                    }
1687                    TapeOp::Neg(a) => {
1688                        adj[*a] -= w;
1689                        adj_dot[*a] -= wd;
1690                    }
1691                    TapeOp::Abs(a) => {
1692                        let s = if v[*a] >= 0.0 { 1.0 } else { -1.0 };
1693                        adj[*a] += w * s;
1694                        adj_dot[*a] += wd * s;
1695                    }
1696                    TapeOp::Sqrt(a) => {
1697                        let sv = v[i];
1698                        if sv > 0.0 {
1699                            let fp = 0.5 / sv;
1700                            let fpp = sqrt_d2(v[*a], sv);
1701                            adj[*a] += w * fp;
1702                            adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1703                        }
1704                    }
1705                    TapeOp::Exp(a) => {
1706                        let ev = v[i];
1707                        adj[*a] += w * ev;
1708                        adj_dot[*a] += wd * ev + w * ev * dot[*a];
1709                    }
1710                    TapeOp::Log(a) => {
1711                        let u = v[*a];
1712                        adj[*a] += w / u;
1713                        adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
1714                    }
1715                    TapeOp::Log10(a) => {
1716                        let u = v[*a];
1717                        adj[*a] += div_u_ln10(w, u);
1718                        adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
1719                    }
1720                    TapeOp::Sin(a) => {
1721                        let u = v[*a];
1722                        let cu = u.cos();
1723                        adj[*a] += w * cu;
1724                        adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1725                    }
1726                    TapeOp::Cos(a) => {
1727                        let u = v[*a];
1728                        let su = u.sin();
1729                        adj[*a] -= w * su;
1730                        adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1731                    }
1732                    TapeOp::Tan(a) => {
1733                        let t = v[i];
1734                        let gp = 1.0 + t * t;
1735                        let gpp = 2.0 * t * gp;
1736                        adj[*a] += w * gp;
1737                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1738                    }
1739                    TapeOp::Atan(a) => {
1740                        let u = v[*a];
1741                        let gp = atan_d1(u);
1742                        let gpp = atan_d2(u);
1743                        adj[*a] += w * gp;
1744                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1745                    }
1746                    TapeOp::Acos(a) => {
1747                        let u = v[*a];
1748                        let gp = acos_d1(u);
1749                        let gpp = acos_d2(u);
1750                        adj[*a] += w * gp;
1751                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1752                    }
1753                    TapeOp::Sinh(a) => {
1754                        let u = v[*a];
1755                        let gp = u.cosh();
1756                        let gpp = u.sinh();
1757                        adj[*a] += w * gp;
1758                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1759                    }
1760                    TapeOp::Cosh(a) => {
1761                        let u = v[*a];
1762                        let gp = u.sinh();
1763                        let gpp = u.cosh();
1764                        adj[*a] += w * gp;
1765                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1766                    }
1767                    TapeOp::Tanh(a) => {
1768                        let t = v[i];
1769                        let gp = tanh_d1(v[*a]);
1770                        let gpp = -2.0 * t * gp;
1771                        adj[*a] += w * gp;
1772                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1773                    }
1774                    TapeOp::Asin(a) => {
1775                        let u = v[*a];
1776                        let gp = asin_d1(u);
1777                        let gpp = asin_d2(u);
1778                        adj[*a] += w * gp;
1779                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1780                    }
1781                    TapeOp::Acosh(a) => {
1782                        let u = v[*a];
1783                        let gp = acosh_d1(u);
1784                        let gpp = acosh_d2(u);
1785                        adj[*a] += w * gp;
1786                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1787                    }
1788                    TapeOp::Asinh(a) => {
1789                        let u = v[*a];
1790                        let gp = asinh_d1(u);
1791                        let gpp = asinh_d2(u);
1792                        adj[*a] += w * gp;
1793                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1794                    }
1795                    TapeOp::Atanh(a) => {
1796                        let u = v[*a];
1797                        let gp = atanh_d1(u);
1798                        let gpp = atanh_d2(u);
1799                        adj[*a] += w * gp;
1800                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1801                    }
1802                    TapeOp::Erf(a) => {
1803                        let u = v[*a];
1804                        let gp = erf_d1(u);
1805                        let gpp = erf_d2(u);
1806                        adj[*a] += w * gp;
1807                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1808                    }
1809                    TapeOp::XLogX(a) => {
1810                        // See the identical arm in `hessian_accumulate`: gpp is
1811                        // 1/u, never the unrepresentable -1/u² of `ln''`.
1812                        let u = v[*a];
1813                        let gp = xlogx_d1(u);
1814                        let gpp = xlogx_d2(u);
1815                        adj[*a] += w * gp;
1816                        adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1817                    }
1818                    TapeOp::CEntropy(a, b) => {
1819                        let ua = v[*a];
1820                        let ub = v[*b];
1821                        let fa = centropy_da(ua, ub);
1822                        let fb = centropy_db(ua, ub);
1823                        let faa = centropy_daa(ua);
1824                        let fab = centropy_dab(ub);
1825                        let fbb = centropy_dbb(ua, ub);
1826                        adj[*a] += w * fa;
1827                        adj[*b] += w * fb;
1828                        adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1829                        adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1830                    }
1831                    TapeOp::Atan2(a, b) => {
1832                        let y = v[*a];
1833                        let x = v[*b];
1834                        let (fa, fb) = atan2_d1(y, x);
1835                        let (faa, fab, fbb) = atan2_d2(y, x);
1836                        adj[*a] += w * fa;
1837                        adj[*b] += w * fb;
1838                        adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1839                        adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1840                    }
1841                    // min/max are piecewise linear (zero second
1842                    // derivative): route adjoint and its tangent into
1843                    // the selected operand, like an active Select branch.
1844                    TapeOp::Min(a, b) => {
1845                        let br = if v[*a] <= v[*b] { *a } else { *b };
1846                        adj[br] += w;
1847                        adj_dot[br] += wd;
1848                    }
1849                    TapeOp::Max(a, b) => {
1850                        let br = if v[*a] >= v[*b] { *a } else { *b };
1851                        adj[br] += w;
1852                        adj_dot[br] += wd;
1853                    }
1854                    // Zero derivative: no first- or second-order adjoint.
1855                    TapeOp::Cmp(_, _, _)
1856                    | TapeOp::And(_, _)
1857                    | TapeOp::Or(_, _)
1858                    | TapeOp::Not(_) => {}
1859                    // Route adjoint and its tangent into the active
1860                    // branch only; the condition contributes nothing.
1861                    TapeOp::Select(c, t, e) => {
1862                        let br = if v[*c] != 0.0 { *t } else { *e };
1863                        adj[br] += w;
1864                        adj_dot[br] += wd;
1865                    }
1866                    TapeOp::Funcall(fc) => {
1867                        let FuncallData { lib, name, args } = fc.as_ref();
1868                        let call_args = funcall_to_ext_args(args, &v);
1869                        let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1870                        let derivs = res.derivs.expect("want_derivs=true returns derivs");
1871                        let hes = res.hessian.expect("want_hes=true returns hessian");
1872                        let real_tape: Vec<usize> = args
1873                            .iter()
1874                            .filter_map(|a| match a {
1875                                TapeFuncallArg::Tape(t) => Some(*t),
1876                                TapeFuncallArg::Str(_) => None,
1877                            })
1878                            .collect();
1879                        for (k, &tk) in real_tape.iter().enumerate() {
1880                            adj[tk] += w * derivs[k];
1881                            let mut second_term = 0.0;
1882                            for (l, &tl) in real_tape.iter().enumerate() {
1883                                let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1884                                let h_kl = hes[lo + hi * (hi + 1) / 2];
1885                                second_term += h_kl * dot[tl];
1886                            }
1887                            adj_dot[tk] += wd * derivs[k] + w * second_term;
1888                        }
1889                    }
1890                }
1891            }
1892        }
1893    }
1894
1895    /// Structural Hessian sparsity (lower triangle, row >= col).
1896    /// Propagates per-slot variable-dependence sets forward; each
1897    /// nonlinear op emits the cross/self products of its operand sets.
1898    /// Linear ops contribute no second-derivative pairs.
1899    pub fn hessian_sparsity(&self) -> BTreeSet<(usize, usize)> {
1900        let n = self.ops.len();
1901        let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
1902        let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
1903
1904        // Slot lifetimes: `last_use[i]` is the highest-numbered op that
1905        // reads slot `i` (or `i` itself when nothing does). A slot's
1906        // variable set is dead the moment its last reader has been
1907        // processed, so it can be *moved into* that reader's result
1908        // instead of copied, and dropped otherwise.
1909        //
1910        // Without this the pass is O(n²) in both time and memory on any
1911        // tape carrying a long additive chain — `v0 * (v0 + v1 + … + vk)`,
1912        // the shape a dense Hessian row comes from — because every
1913        // partial sum leaves behind its own full-size `BTreeSet` and each
1914        // union copies the whole running set. `split_top_sums` keeps
1915        // *top-level* sums out of a single tape, so this only bites when
1916        // the sum feeds an enclosing operator, which is exactly the dense
1917        // case. With reuse, a left-leaning chain extends one set in place.
1918        let mut last_use: Vec<usize> = (0..n).collect();
1919        for (i, op) in self.ops.iter().enumerate() {
1920            for_each_input(op, |a| last_use[a] = i);
1921        }
1922        // Union of two operand sets, taking ownership of whichever dies
1923        // at this op rather than copying. `a == b` is handled first:
1924        // taking one would empty the other.
1925        macro_rules! merge {
1926            ($a:expr, $b:expr, $i:expr) => {{
1927                let (a, b, i) = ($a, $b, $i);
1928                if a == b {
1929                    if last_use[a] == i {
1930                        std::mem::take(&mut var_sets[a])
1931                    } else {
1932                        var_sets[a].clone()
1933                    }
1934                } else {
1935                    // Prefer moving the larger of the two dead sets.
1936                    let take_a = last_use[a] == i
1937                        && (last_use[b] != i || var_sets[a].len() >= var_sets[b].len());
1938                    if take_a {
1939                        let mut s = std::mem::take(&mut var_sets[a]);
1940                        s.extend(var_sets[b].iter().copied());
1941                        s
1942                    } else if last_use[b] == i {
1943                        let mut s = std::mem::take(&mut var_sets[b]);
1944                        s.extend(var_sets[a].iter().copied());
1945                        s
1946                    } else {
1947                        var_sets[a].union(&var_sets[b]).copied().collect()
1948                    }
1949                }
1950            }};
1951        }
1952        // Same, for the one-operand ops.
1953        macro_rules! carry {
1954            ($a:expr, $i:expr) => {{
1955                let (a, i) = ($a, $i);
1956                if last_use[a] == i {
1957                    std::mem::take(&mut var_sets[a])
1958                } else {
1959                    var_sets[a].clone()
1960                }
1961            }};
1962        }
1963
1964        let emit_cross =
1965            |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1966                for &v1 in s1 {
1967                    for &v2 in s2 {
1968                        let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
1969                        pairs.insert((r, c));
1970                    }
1971                }
1972            };
1973        let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1974            let vars: Vec<usize> = s.iter().copied().collect();
1975            for (ai, &vi) in vars.iter().enumerate() {
1976                for &vj in &vars[..=ai] {
1977                    let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
1978                    pairs.insert((r, c));
1979                }
1980            }
1981        };
1982
1983        for (i, op) in self.ops.iter().enumerate() {
1984            let vset = match op {
1985                TapeOp::Const(_) => BTreeSet::new(),
1986                TapeOp::Var(j) => {
1987                    let mut s = BTreeSet::new();
1988                    s.insert(*j);
1989                    s
1990                }
1991                TapeOp::Add(a, b) | TapeOp::Sub(a, b) => merge!(*a, *b, i),
1992                TapeOp::Neg(a) | TapeOp::Abs(a) => carry!(*a, i),
1993                TapeOp::Mul(a, b) => {
1994                    emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1995                    merge!(*a, *b, i)
1996                }
1997                TapeOp::Div(a, b) => {
1998                    emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1999                    emit_self(&var_sets[*b], &mut pairs);
2000                    merge!(*a, *b, i)
2001                }
2002                TapeOp::Pow(a, b) => {
2003                    let combined = merge!(*a, *b, i);
2004                    emit_self(&combined, &mut pairs);
2005                    combined
2006                }
2007                TapeOp::Sqrt(a)
2008                | TapeOp::Exp(a)
2009                | TapeOp::Log(a)
2010                | TapeOp::Log10(a)
2011                | TapeOp::Sin(a)
2012                | TapeOp::Cos(a)
2013                | TapeOp::Tan(a)
2014                | TapeOp::Atan(a)
2015                | TapeOp::Acos(a)
2016                | TapeOp::Sinh(a)
2017                | TapeOp::Cosh(a)
2018                | TapeOp::Tanh(a)
2019                | TapeOp::Asin(a)
2020                | TapeOp::Acosh(a)
2021                | TapeOp::Asinh(a)
2022                | TapeOp::Erf(a)
2023                | TapeOp::XLogX(a)
2024                | TapeOp::Atanh(a) => {
2025                    emit_self(&var_sets[*a], &mut pairs);
2026                    carry!(*a, i)
2027                }
2028                // atan2(y, x) and centropy(a, b) are nonlinear in both
2029                // operands with a full 2×2 second-derivative block; the
2030                // structural superset is every self/cross pair within the
2031                // combined operand set.
2032                TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
2033                    let combined = merge!(*a, *b, i);
2034                    emit_self(&combined, &mut pairs);
2035                    combined
2036                }
2037                // Comparisons / logical connectives are piecewise
2038                // constant: identically-zero derivative, so they
2039                // introduce no second-derivative pairs and carry no
2040                // variable dependence downstream (their result is a
2041                // constant as far as AD is concerned).
2042                TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
2043                    BTreeSet::new()
2044                }
2045                // Select passes through the active branch's value with
2046                // unit derivative, so it emits no pairs of its own; its
2047                // dependence set is the union of *both* branches
2048                // (either may become active as x varies — conservative
2049                // and correct for a structural superset). The condition
2050                // contributes no derivative and is excluded.
2051                TapeOp::Select(_c, t, e) => merge!(*t, *e, i),
2052                // min/max are piecewise linear: the active operand passes
2053                // through with unit derivative, so the second derivative is
2054                // identically zero (no pairs). Their dependence set is the
2055                // union of both operands (either may become active as x
2056                // varies — conservative and correct for a structural
2057                // superset), mirroring Select.
2058                TapeOp::Min(a, b) | TapeOp::Max(a, b) => merge!(*a, *b, i),
2059                TapeOp::Funcall(fc) => {
2060                    let args = &fc.args;
2061                    let mut combined: BTreeSet<usize> = BTreeSet::new();
2062                    for arg in args {
2063                        if let TapeFuncallArg::Tape(t) = arg {
2064                            for &vv in &var_sets[*t] {
2065                                combined.insert(vv);
2066                            }
2067                        }
2068                    }
2069                    emit_self(&combined, &mut pairs);
2070                    combined
2071                }
2072            };
2073            // Release any operand whose last reader was this op and which
2074            // `merge!` / `carry!` did not already move out (already-moved
2075            // slots are empty, so this is a no-op for them). Without it a
2076            // tape's peak memory is the sum of every intermediate set
2077            // rather than the live ones.
2078            for_each_input(op, |a| {
2079                if last_use[a] == i {
2080                    var_sets[a].clear();
2081                }
2082            });
2083            var_sets.push(vset);
2084        }
2085        pairs
2086    }
2087}
2088
2089/// Call `f` on every tape slot `op` reads.
2090///
2091/// Deliberately written as an exhaustive match with **no catch-all arm**:
2092/// a new `TapeOp` variant must be added here or this stops compiling.
2093/// A silent "reads nothing" default would let `hessian_sparsity`'s
2094/// liveness analysis drop an operand's variable set while it is still
2095/// needed, which would quietly under-report Hessian sparsity — the kind
2096/// of wrong answer that surfaces as a bad search direction, not a panic.
2097/// Over-reporting an input is merely a missed optimization, so when in
2098/// doubt list the slot.
2099fn for_each_input(op: &TapeOp, mut f: impl FnMut(usize)) {
2100    match op {
2101        TapeOp::Const(_) | TapeOp::Var(_) => {}
2102        TapeOp::Neg(a)
2103        | TapeOp::Abs(a)
2104        | TapeOp::Sqrt(a)
2105        | TapeOp::Exp(a)
2106        | TapeOp::Log(a)
2107        | TapeOp::Log10(a)
2108        | TapeOp::Sin(a)
2109        | TapeOp::Cos(a)
2110        | TapeOp::Tan(a)
2111        | TapeOp::Atan(a)
2112        | TapeOp::Acos(a)
2113        | TapeOp::Sinh(a)
2114        | TapeOp::Cosh(a)
2115        | TapeOp::Tanh(a)
2116        | TapeOp::Asin(a)
2117        | TapeOp::Acosh(a)
2118        | TapeOp::Asinh(a)
2119        | TapeOp::Atanh(a)
2120        | TapeOp::Erf(a)
2121        | TapeOp::XLogX(a)
2122        | TapeOp::Not(a) => f(*a),
2123        TapeOp::Add(a, b)
2124        | TapeOp::Sub(a, b)
2125        | TapeOp::Mul(a, b)
2126        | TapeOp::Div(a, b)
2127        | TapeOp::Pow(a, b)
2128        | TapeOp::CEntropy(a, b)
2129        | TapeOp::Atan2(a, b)
2130        | TapeOp::Min(a, b)
2131        | TapeOp::Max(a, b)
2132        | TapeOp::And(a, b)
2133        | TapeOp::Or(a, b)
2134        | TapeOp::Cmp(_, a, b) => {
2135            f(*a);
2136            f(*b);
2137        }
2138        TapeOp::Select(c, t, e) => {
2139            f(*c);
2140            f(*t);
2141            f(*e);
2142        }
2143        TapeOp::Funcall(fc) => {
2144            for arg in &fc.args {
2145                if let TapeFuncallArg::Tape(t) = arg {
2146                    f(*t);
2147                }
2148            }
2149        }
2150    }
2151}
2152
2153fn build_recursive(
2154    expr: &Expr,
2155    ops: &mut Vec<TapeOp>,
2156    cache: &mut HashMap<*const Expr, usize>,
2157    resolver: &ExternalResolver,
2158) -> usize {
2159    match expr {
2160        Expr::Const(c) => {
2161            let idx = ops.len();
2162            ops.push(TapeOp::Const(*c));
2163            idx
2164        }
2165        Expr::Var(i) => {
2166            let idx = ops.len();
2167            ops.push(TapeOp::Var(*i));
2168            idx
2169        }
2170        Expr::Binary(op, a, b) => {
2171            // Pow(x, const) is the dominant libm/dispatch cost in
2172            // transcendental-heavy AMPL tapes (henon, lane_emden, …):
2173            // `powf` itself is ~30–50 cycles AND the reverse-mode arm
2174            // for `Pow` carries an extra `ln(x)` branch. Rewriting
2175            // small integer / half-integer exponents into mul/sqrt
2176            // chains drops these calls entirely and reroutes the AD
2177            // through the much cheaper `Mul`/`Sqrt` arms.
2178            if let BinOp::Pow = op {
2179                if let Some(c) = peek_const(b) {
2180                    if let Some(idx) = try_emit_const_pow(a, c, ops, cache, resolver) {
2181                        return idx;
2182                    }
2183                }
2184            }
2185            let l = build_recursive(a, ops, cache, resolver);
2186            let r = build_recursive(b, ops, cache, resolver);
2187            let idx = ops.len();
2188            ops.push(match op {
2189                BinOp::Add => TapeOp::Add(l, r),
2190                BinOp::Sub => TapeOp::Sub(l, r),
2191                BinOp::Mul => TapeOp::Mul(l, r),
2192                BinOp::Div => TapeOp::Div(l, r),
2193                BinOp::Pow => TapeOp::Pow(l, r),
2194                BinOp::Atan2 => TapeOp::Atan2(l, r),
2195                BinOp::CEntropy => TapeOp::CEntropy(l, r),
2196            });
2197            idx
2198        }
2199        Expr::Unary(op, a) => {
2200            let v = build_recursive(a, ops, cache, resolver);
2201            let idx = ops.len();
2202            ops.push(match op {
2203                UnaryOp::Neg => TapeOp::Neg(v),
2204                UnaryOp::Sqrt => TapeOp::Sqrt(v),
2205                UnaryOp::Log => TapeOp::Log(v),
2206                UnaryOp::Log10 => TapeOp::Log10(v),
2207                UnaryOp::Exp => TapeOp::Exp(v),
2208                UnaryOp::Abs => TapeOp::Abs(v),
2209                UnaryOp::Sin => TapeOp::Sin(v),
2210                UnaryOp::Cos => TapeOp::Cos(v),
2211                UnaryOp::Tan => TapeOp::Tan(v),
2212                UnaryOp::Atan => TapeOp::Atan(v),
2213                UnaryOp::Acos => TapeOp::Acos(v),
2214                UnaryOp::Sinh => TapeOp::Sinh(v),
2215                UnaryOp::Cosh => TapeOp::Cosh(v),
2216                UnaryOp::Tanh => TapeOp::Tanh(v),
2217                UnaryOp::Asin => TapeOp::Asin(v),
2218                UnaryOp::Acosh => TapeOp::Acosh(v),
2219                UnaryOp::Asinh => TapeOp::Asinh(v),
2220                UnaryOp::Atanh => TapeOp::Atanh(v),
2221                UnaryOp::Erf => TapeOp::Erf(v),
2222                UnaryOp::XLogX => TapeOp::XLogX(v),
2223            });
2224            idx
2225        }
2226        Expr::Sum(args) => {
2227            if args.is_empty() {
2228                let idx = ops.len();
2229                ops.push(TapeOp::Const(0.0));
2230                return idx;
2231            }
2232            let mut acc = build_recursive(&args[0], ops, cache, resolver);
2233            for a in &args[1..] {
2234                let next = build_recursive(a, ops, cache, resolver);
2235                let idx = ops.len();
2236                ops.push(TapeOp::Add(acc, next));
2237                acc = idx;
2238            }
2239            acc
2240        }
2241        // n-ary min/max fold to a left-associative chain of binary
2242        // Min/Max TapeOps. The chain reproduces the list extremum, and
2243        // the binary Min/Max AD arms route the (sub)gradient to the
2244        // active operand at each step — equivalent to selecting the one
2245        // active operand of the whole list. An empty list cannot arise
2246        // from a well-formed `.nl` MINLIST/MAXLIST (count >= 1); guard
2247        // with a 0 constant for safety rather than panicking.
2248        Expr::MinList(args) | Expr::MaxList(args) => {
2249            let is_min = matches!(expr, Expr::MinList(_));
2250            if args.is_empty() {
2251                let idx = ops.len();
2252                ops.push(TapeOp::Const(0.0));
2253                return idx;
2254            }
2255            let mut acc = build_recursive(&args[0], ops, cache, resolver);
2256            for a in &args[1..] {
2257                let next = build_recursive(a, ops, cache, resolver);
2258                let idx = ops.len();
2259                ops.push(if is_min {
2260                    TapeOp::Min(acc, next)
2261                } else {
2262                    TapeOp::Max(acc, next)
2263                });
2264                acc = idx;
2265            }
2266            acc
2267        }
2268        Expr::Cse(body) => {
2269            // Cache by Arc identity so each shared body is emitted into
2270            // the tape exactly once and every reference resolves to the
2271            // same result-slot index. Forward computes the body once;
2272            // reverse-mode adjoint sums contributions from every ref
2273            // into that shared slot — exact chain rule for shared
2274            // sub-expressions.
2275            let key = Arc::as_ptr(body) as *const Expr;
2276            if let Some(&idx) = cache.get(&key) {
2277                idx
2278            } else {
2279                let idx = build_recursive(body, ops, cache, resolver);
2280                cache.insert(key, idx);
2281                idx
2282            }
2283        }
2284        Expr::Compare(op, a, b) => {
2285            let l = build_recursive(a, ops, cache, resolver);
2286            let r = build_recursive(b, ops, cache, resolver);
2287            let idx = ops.len();
2288            ops.push(TapeOp::Cmp(*op, l, r));
2289            idx
2290        }
2291        Expr::And(a, b) => {
2292            let l = build_recursive(a, ops, cache, resolver);
2293            let r = build_recursive(b, ops, cache, resolver);
2294            let idx = ops.len();
2295            ops.push(TapeOp::And(l, r));
2296            idx
2297        }
2298        Expr::Or(a, b) => {
2299            let l = build_recursive(a, ops, cache, resolver);
2300            let r = build_recursive(b, ops, cache, resolver);
2301            let idx = ops.len();
2302            ops.push(TapeOp::Or(l, r));
2303            idx
2304        }
2305        Expr::Not(a) => {
2306            let v = build_recursive(a, ops, cache, resolver);
2307            let idx = ops.len();
2308            ops.push(TapeOp::Not(v));
2309            idx
2310        }
2311        Expr::Cond { cond, then_, else_ } => {
2312            let c = build_recursive(cond, ops, cache, resolver);
2313            let t = build_recursive(then_, ops, cache, resolver);
2314            let e = build_recursive(else_, ops, cache, resolver);
2315            let idx = ops.len();
2316            ops.push(TapeOp::Select(c, t, e));
2317            idx
2318        }
2319        Expr::Funcall { id, args } => {
2320            let (lib, name) = resolver
2321                .funcs_by_id
2322                .get(id)
2323                .unwrap_or_else(|| panic!("unresolved AMPL funcall id {id}"));
2324            let tape_args: Vec<TapeFuncallArg> = args
2325                .iter()
2326                .map(|a| match a {
2327                    FuncallArg::Real(e) => {
2328                        TapeFuncallArg::Tape(build_recursive(e, ops, cache, resolver))
2329                    }
2330                    FuncallArg::Str(s) => TapeFuncallArg::Str(s.clone()),
2331                })
2332                .collect();
2333            let idx = ops.len();
2334            ops.push(TapeOp::Funcall(Box::new(FuncallData {
2335                lib: Arc::clone(lib),
2336                name: name.clone(),
2337                args: tape_args,
2338            })));
2339            idx
2340        }
2341    }
2342}
2343
2344/// Resolve `e` to a literal constant if it is one (transparently
2345/// peering through `Cse` wrappers, which AMPL emits around shared
2346/// constants in CSE-heavy problems).
2347fn peek_const(e: &Expr) -> Option<f64> {
2348    match e {
2349        Expr::Const(c) => Some(*c),
2350        Expr::Cse(body) => peek_const(body),
2351        _ => None,
2352    }
2353}
2354
2355/// Try to rewrite `base ^ exponent_const` into cheaper ops. Returns
2356/// the result tape-slot on success; `None` means "fall through to
2357/// generic Pow." Handles the cases that account for the bulk of
2358/// AMPL-emitted Pow nodes: integer exponents up to ±8 and the
2359/// `Sqrt`/passthrough/one specials. Half-integer exponents (e.g.
2360/// `^1.5`) and larger integers are left to generic `Pow` since the
2361/// resulting mul chain grows the tape faster than it saves work.
2362fn try_emit_const_pow(
2363    base_expr: &Expr,
2364    c: f64,
2365    ops: &mut Vec<TapeOp>,
2366    cache: &mut HashMap<*const Expr, usize>,
2367    resolver: &ExternalResolver,
2368) -> Option<usize> {
2369    if c == 0.0 {
2370        let idx = ops.len();
2371        ops.push(TapeOp::Const(1.0));
2372        return Some(idx);
2373    }
2374    if c == 1.0 {
2375        return Some(build_recursive(base_expr, ops, cache, resolver));
2376    }
2377    if c == 0.5 {
2378        let b = build_recursive(base_expr, ops, cache, resolver);
2379        let idx = ops.len();
2380        ops.push(TapeOp::Sqrt(b));
2381        return Some(idx);
2382    }
2383    // Integer exponents: bounded so a bad tape can't blow up the
2384    // op count. 8 covers everything AMPL typically emits for
2385    // polynomial models; beyond that the binary-expansion mul
2386    // chain (≥4 ops) starts to lose to a single `powf`.
2387    if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
2388        let n = c.abs() as u32;
2389        if n == 0 {
2390            // Already handled above, but guard.
2391            let idx = ops.len();
2392            ops.push(TapeOp::Const(1.0));
2393            return Some(idx);
2394        }
2395        let b = build_recursive(base_expr, ops, cache, resolver);
2396        let pos = emit_int_pow(b, n, ops);
2397        if c < 0.0 {
2398            // x^-n = 1 / x^n. Saves the powf and its ln branch in
2399            // reverse mode; cost is one Div in their place.
2400            let one_idx = ops.len();
2401            ops.push(TapeOp::Const(1.0));
2402            let idx = ops.len();
2403            ops.push(TapeOp::Div(one_idx, pos));
2404            return Some(idx);
2405        }
2406        return Some(pos);
2407    }
2408    None
2409}
2410
2411/// Emit `base^n` for `n >= 1` as a binary-expansion mul chain.
2412/// Worst-case op count is `2·floor(log2(n))` — i.e. 1 op for n=2, 2
2413/// for n=3/4, 3 for n=5..8.
2414fn emit_int_pow(base: usize, n: u32, ops: &mut Vec<TapeOp>) -> usize {
2415    debug_assert!(n >= 1);
2416    if n == 1 {
2417        return base;
2418    }
2419    let half = emit_int_pow(base, n / 2, ops);
2420    let squared = ops.len();
2421    ops.push(TapeOp::Mul(half, half));
2422    if n % 2 == 1 {
2423        let idx = ops.len();
2424        ops.push(TapeOp::Mul(squared, base));
2425        idx
2426    } else {
2427        squared
2428    }
2429}
2430
2431// ============================================================
2432// HybridTape: per-summand local tapes + shared CSE prelude.
2433//
2434// Partial separability — the .nl Sum/Add structure — gets each
2435// summand its own local Vec<SummandOp>. CSE bodies (V-segments
2436// in .nl) that appear in two or more summands are promoted into
2437// a single shared `prelude: Vec<TapeOp>`; per-summand references
2438// to a promoted CSE are SummandOp::Shared(prelude_slot).
2439//
2440// This is strictly better than either extreme:
2441//   - per-summand Tape (no cross-summand sharing): re-inlines
2442//     every shared CSE, blows up tape size when many constraints
2443//     share a stencil derivative (Mittelmann *120 problems).
2444//   - GlobalTape (single shared Vec<TapeOp> for everything):
2445//     per-root reverse sweeps scatter across a many-MB buffer,
2446//     thrashing cache when no CSE is actually shared (lane_emden
2447//     120: each constraint owns its own ops → 50% regression
2448//     vs per-summand tapes).
2449//
2450// Forward: prelude once, then each summand's local pass.
2451// Reverse / forward-over-reverse: per-summand sweep over local
2452// reach (which propagates adjoints into prelude_adj at Shared
2453// boundaries), then a small reverse pass over the summand's
2454// prelude_reach to fold those into grad / Hessian.
2455// ============================================================
2456
2457/// One slot in a per-summand local tape.
2458#[derive(Debug, Clone)]
2459pub enum SummandOp {
2460    /// Local op — operand indices reference other slots in the
2461    /// same per-summand vector.
2462    Local(TapeOp),
2463    /// Pull a value from the shared prelude at slot `usize`. No
2464    /// downstream cost beyond the lookup; adjoints flowing into
2465    /// this slot accumulate into the prelude adjoint buffer.
2466    Shared(usize),
2467}
2468
2469#[derive(Debug, Clone)]
2470pub struct Summand {
2471    pub ops: Vec<SummandOp>,
2472    /// Local slot holding the summand's final value.
2473    pub root_slot: usize,
2474    /// Local slots reachable from `root_slot`, ascending (topo).
2475    pub local_reach: Vec<usize>,
2476    /// Prelude slots reachable from the summand's Shared refs,
2477    /// ascending (topo in prelude's operand DAG).
2478    pub prelude_reach: Vec<usize>,
2479    /// Variables touched by Var ops inside `local_reach`.
2480    pub local_vars: Vec<usize>,
2481    /// Variables touched by Var ops inside `prelude_reach`.
2482    pub prelude_vars: Vec<usize>,
2483    /// `local_vars ∪ prelude_vars`, sorted — every problem variable this
2484    /// summand can contribute a derivative for. `eval_h` maps it through
2485    /// the Hessian coloring to decide which colors the summand
2486    /// participates in (issue #557); it must include `prelude_vars` or a
2487    /// summand would be skipped for colors it reaches only through a
2488    /// shared CSE body.
2489    pub all_vars: Vec<usize>,
2490}
2491
2492#[derive(Debug, Clone)]
2493pub struct HybridTape {
2494    /// Shared CSE bodies. Slot indices in `SummandOp::Shared`
2495    /// point here; this Vec is built bottom-up by `build_recursive`,
2496    /// so operand indices are always less than the consumer's
2497    /// index (topo in ascending order).
2498    pub prelude: Vec<TapeOp>,
2499    pub summands: Vec<Summand>,
2500}
2501
2502impl HybridTape {
2503    /// Build hybrid tape from a list of root expressions. CSE
2504    /// bodies referenced from ≥ 2 roots are promoted into the
2505    /// shared prelude; CSEs touched by only one root are inlined
2506    /// into that summand's local ops.
2507    pub fn build_multi(exprs: &[Expr]) -> Self {
2508        // Pass 1: per-Cse-pointer count of how many roots reference
2509        // it (each root contributes at most 1 to the count). The
2510        // ≥2 threshold means a CSE is shared across summands.
2511        let mut cse_count: HashMap<*const Expr, usize> = HashMap::new();
2512        for e in exprs {
2513            let mut seen_in_root: HashSet<*const Expr> = HashSet::new();
2514            count_cse_appearances(e, &mut seen_in_root, &mut cse_count);
2515        }
2516
2517        // Pass 2: build prelude + each summand. The summand builder
2518        // hits the prelude path lazily — only when it encounters a
2519        // promoted Cse — so the prelude grows only with bodies that
2520        // are actually referenced from multiple summands.
2521        let mut prelude: Vec<TapeOp> = Vec::new();
2522        let mut prelude_map: HashMap<*const Expr, usize> = HashMap::new();
2523        let mut summands: Vec<Summand> = Vec::with_capacity(exprs.len());
2524        for e in exprs {
2525            let mut local: Vec<SummandOp> = Vec::new();
2526            let mut local_cache: HashMap<*const Expr, usize> = HashMap::new();
2527            let root_slot = build_into_summand(
2528                e,
2529                &mut local,
2530                &mut local_cache,
2531                &mut prelude,
2532                &mut prelude_map,
2533                &cse_count,
2534            );
2535            summands.push(Summand {
2536                ops: local,
2537                root_slot,
2538                local_reach: Vec::new(),
2539                prelude_reach: Vec::new(),
2540                local_vars: Vec::new(),
2541                prelude_vars: Vec::new(),
2542                all_vars: Vec::new(),
2543            });
2544        }
2545
2546        // Pass 3: per-summand reach / vars. Prelude reach uses an
2547        // epoch-tagged shared visited buffer so total cost stays
2548        // O(Σ |prelude_reach_i|) rather than O(n_summands × |prelude|).
2549        let mut p_visited: Vec<u32> = vec![0; prelude.len()];
2550        let mut p_epoch: u32 = 0;
2551        let mut p_stack: Vec<usize> = Vec::new();
2552        for s in &mut summands {
2553            let (local_reach, shared_refs) = compute_local_reach(&s.ops, s.root_slot);
2554            s.local_reach = local_reach;
2555
2556            let mut lv: BTreeSet<usize> = BTreeSet::new();
2557            for &i in &s.local_reach {
2558                if let SummandOp::Local(TapeOp::Var(j)) = &s.ops[i] {
2559                    lv.insert(*j);
2560                }
2561            }
2562            s.local_vars = lv.iter().copied().collect();
2563
2564            if !shared_refs.is_empty() {
2565                p_epoch += 1;
2566                let mut preach: Vec<usize> = Vec::new();
2567                for &start in &shared_refs {
2568                    bfs_prelude(
2569                        &prelude,
2570                        start,
2571                        &mut p_visited,
2572                        p_epoch,
2573                        &mut p_stack,
2574                        &mut preach,
2575                    );
2576                }
2577                preach.sort_unstable();
2578                s.prelude_vars = vars_in(&prelude, &preach);
2579                s.prelude_reach = preach;
2580            }
2581
2582            let mut av: BTreeSet<usize> = lv;
2583            for &v in &s.prelude_vars {
2584                av.insert(v);
2585            }
2586            s.all_vars = av.into_iter().collect();
2587        }
2588
2589        HybridTape { prelude, summands }
2590    }
2591
2592    pub fn n_prelude_ops(&self) -> usize {
2593        self.prelude.len()
2594    }
2595    pub fn n_summands(&self) -> usize {
2596        self.summands.len()
2597    }
2598    pub fn max_summand_ops(&self) -> usize {
2599        self.summands.iter().map(|s| s.ops.len()).max().unwrap_or(0)
2600    }
2601    pub fn total_local_ops(&self) -> usize {
2602        self.summands.iter().map(|s| s.ops.len()).sum()
2603    }
2604
2605    /// Forward sweep over the shared prelude. `prelude_vals` must
2606    /// have length `n_prelude_ops`.
2607    pub fn forward_prelude(&self, x: &[f64], prelude_vals: &mut [f64]) {
2608        debug_assert_eq!(prelude_vals.len(), self.prelude.len());
2609        for i in 0..self.prelude.len() {
2610            prelude_vals[i] = fwd_step(&self.prelude[i], x, prelude_vals);
2611        }
2612    }
2613
2614    /// Forward sweep over one summand. `local_vals` must hold at
2615    /// least `s.ops.len()` entries.
2616    pub fn forward_summand(
2617        &self,
2618        s: &Summand,
2619        x: &[f64],
2620        prelude_vals: &[f64],
2621        local_vals: &mut [f64],
2622    ) {
2623        debug_assert!(local_vals.len() >= s.ops.len());
2624        for i in 0..s.ops.len() {
2625            local_vals[i] = match &s.ops[i] {
2626                SummandOp::Local(op) => fwd_step(op, x, local_vals),
2627                SummandOp::Shared(k) => prelude_vals[*k],
2628            };
2629        }
2630    }
2631
2632    /// Value at the summand root after `forward_summand`.
2633    #[inline]
2634    pub fn root_value(&self, s: &Summand, local_vals: &[f64]) -> f64 {
2635        local_vals[s.root_slot]
2636    }
2637
2638    /// Reverse-mode gradient for one summand. Walks `local_reach`
2639    /// in reverse — propagating adjoints into `prelude_adj` at
2640    /// Shared boundaries — and then walks `prelude_reach` in
2641    /// reverse to land contributions in `grad`. Scratch arrays
2642    /// `local_adj` and `prelude_adj` are zeroed only at the slots
2643    /// actually touched.
2644    #[allow(clippy::too_many_arguments)]
2645    pub fn gradient_summand(
2646        &self,
2647        s: &Summand,
2648        prelude_vals: &[f64],
2649        local_vals: &[f64],
2650        seed: f64,
2651        grad: &mut [f64],
2652        local_adj: &mut [f64],
2653        prelude_adj: &mut [f64],
2654    ) {
2655        if seed == 0.0 || s.local_reach.is_empty() {
2656            return;
2657        }
2658        for &i in &s.local_reach {
2659            local_adj[i] = 0.0;
2660        }
2661        for &i in &s.prelude_reach {
2662            prelude_adj[i] = 0.0;
2663        }
2664        local_adj[s.root_slot] = seed;
2665        for &i in s.local_reach.iter().rev() {
2666            let a = local_adj[i];
2667            if a == 0.0 {
2668                continue;
2669            }
2670            match &s.ops[i] {
2671                SummandOp::Local(op) => rev_step(op, i, local_vals, local_adj, a, grad),
2672                SummandOp::Shared(k) => {
2673                    prelude_adj[*k] += a;
2674                }
2675            }
2676        }
2677        for &i in s.prelude_reach.iter().rev() {
2678            let a = prelude_adj[i];
2679            if a == 0.0 {
2680                continue;
2681            }
2682            rev_step(&self.prelude[i], i, prelude_vals, prelude_adj, a, grad);
2683        }
2684    }
2685
2686    /// Forward tangent over the whole prelude for one dense seed
2687    /// vector (a Hessian color): `prelude_dot[i] = (∂prelude[i]/∂x) ·
2688    /// seed`. Runs **once per color for the entire constraint block**
2689    /// — this is the sweep the coloring lets every summand of that
2690    /// color share, and the reason the directional Hessian path
2691    /// exists at all (issue #557): a per-summand-per-variable seeding
2692    /// strategy would repeat it for every referencing summand.
2693    /// `reach` is the set of prelude slots this color actually uses —
2694    /// the union of `prelude_reach` over the color's summands. It must
2695    /// be **ascending** (the prelude is emitted bottom-up, so operand
2696    /// indices are always below their consumer's) and **closed under
2697    /// operands**, which a union of `prelude_reach` sets is by
2698    /// construction. Slots outside it are neither written nor read:
2699    /// a summand only pulls `prelude_dot[k]` for `k` in its own
2700    /// `prelude_reach`, so a stale tangent left in an untouched slot by
2701    /// a previous color can never be observed. Passing every slot is
2702    /// therefore always correct, just not always cheap — walking the
2703    /// whole prelude once per color is `n_colors × |prelude|` work
2704    /// against the `|prelude|` the op-ratio gate assumes.
2705    pub fn prelude_tangent(
2706        &self,
2707        prelude_vals: &[f64],
2708        seed: &[f64],
2709        reach: &[u32],
2710        prelude_dot: &mut [f64],
2711    ) {
2712        debug_assert!(prelude_dot.len() >= self.prelude.len());
2713        for &i in reach {
2714            let i = i as usize;
2715            prelude_dot[i] = fwd_dir_step(&self.prelude[i], seed, prelude_vals, prelude_dot, i);
2716        }
2717    }
2718
2719    /// Directional forward-over-reverse Hessian pass for one summand
2720    /// with multiplier `weight`, sharing the per-color prelude tangent
2721    /// computed by [`prelude_tangent`]: forward-tangent the local ops
2722    /// (pulling `prelude_dot` at `Shared` boundaries), then
2723    /// reverse-over-tangent them, writing `weight * (H · seed)[k]`
2724    /// into the dense `out[k]` at each `Var(k)` — the same contract as
2725    /// [`Tape::hessian_directional`], so contributions land straight
2726    /// in the coloring's `compressed[c]` buffer.
2727    ///
2728    /// Adjoints crossing a `Shared` boundary are folded into
2729    /// `prelude_adj` / `prelude_adj_dot` **scaled by `weight`**, and
2730    /// left there: reverse-over-tangent is linear in its adjoint
2731    /// seeds, so accumulating `λ_k`-weighted adjoints across every
2732    /// summand of a color and running [`prelude_reverse_directional`]
2733    /// once with unit weight computes the same second-order prelude
2734    /// contribution as a per-summand sweep — while paying for the
2735    /// prelude walk once per color instead of once per (summand,
2736    /// color) pair.
2737    ///
2738    /// `local_vals` must hold this summand's forward values (see
2739    /// [`forward_summand`]); `local_dot` / `local_adj` /
2740    /// `local_adj_dot` are scratch arenas (≥ `s.ops.len()`), zeroed
2741    /// here only at the slots the summand reaches.
2742    ///
2743    /// [`prelude_tangent`]: HybridTape::prelude_tangent
2744    /// [`prelude_reverse_directional`]: HybridTape::prelude_reverse_directional
2745    /// [`forward_summand`]: HybridTape::forward_summand
2746    #[allow(clippy::too_many_arguments)]
2747    pub fn hessian_summand_directional(
2748        &self,
2749        s: &Summand,
2750        local_vals: &[f64],
2751        prelude_dot: &[f64],
2752        seed: &[f64],
2753        weight: f64,
2754        out: &mut [f64],
2755        local_dot: &mut [f64],
2756        local_adj: &mut [f64],
2757        local_adj_dot: &mut [f64],
2758        prelude_adj: &mut [f64],
2759        prelude_adj_dot: &mut [f64],
2760    ) {
2761        if weight == 0.0 || s.local_reach.is_empty() {
2762            return;
2763        }
2764        for &i in &s.local_reach {
2765            local_adj[i] = 0.0;
2766            local_adj_dot[i] = 0.0;
2767        }
2768        // `local_dot` needs no pre-zeroing: `local_reach` is ascending
2769        // and every operand of a reached op is itself reached, so each
2770        // slot is written before any read.
2771        for &i in &s.local_reach {
2772            local_dot[i] = match &s.ops[i] {
2773                SummandOp::Local(op) => fwd_dir_step(op, seed, local_vals, local_dot, i),
2774                SummandOp::Shared(k) => prelude_dot[*k],
2775            };
2776        }
2777        local_adj[s.root_slot] = 1.0;
2778        for &i in s.local_reach.iter().rev() {
2779            let w = local_adj[i];
2780            let wd = local_adj_dot[i];
2781            if w == 0.0 && wd == 0.0 {
2782                continue;
2783            }
2784            match &s.ops[i] {
2785                SummandOp::Local(op) => {
2786                    ror_dir_step(
2787                        op,
2788                        i,
2789                        local_vals,
2790                        local_dot,
2791                        local_adj,
2792                        local_adj_dot,
2793                        w,
2794                        wd,
2795                        weight,
2796                        out,
2797                    );
2798                }
2799                SummandOp::Shared(k) => {
2800                    prelude_adj[*k] += weight * w;
2801                    prelude_adj_dot[*k] += weight * wd;
2802                }
2803            }
2804        }
2805    }
2806
2807    /// Reverse-over-tangent over the prelude with unit weight,
2808    /// consuming the adjoint accumulators built by
2809    /// [`hessian_summand_directional`] across all summands of one
2810    /// color and writing variable contributions into the dense `out`.
2811    ///
2812    /// `prelude_adj` / `prelude_adj_dot` must be all-zero except for
2813    /// the accumulated seeds on entry (allocate them zeroed and this
2814    /// invariant maintains itself); each slot is re-zeroed as it is
2815    /// consumed — adjoints only ever propagate to lower slots, which
2816    /// the reverse walk has not visited yet — so the buffers come back
2817    /// all-zero, ready for the next color, without an O(prelude) fill.
2818    ///
2819    /// [`hessian_summand_directional`]: HybridTape::hessian_summand_directional
2820    pub fn prelude_reverse_directional(
2821        &self,
2822        prelude_vals: &[f64],
2823        prelude_dot: &[f64],
2824        reach: &[u32],
2825        out: &mut [f64],
2826        prelude_adj: &mut [f64],
2827        prelude_adj_dot: &mut [f64],
2828    ) {
2829        for &i in reach.iter().rev() {
2830            let i = i as usize;
2831            let w = prelude_adj[i];
2832            let wd = prelude_adj_dot[i];
2833            if w == 0.0 && wd == 0.0 {
2834                continue;
2835            }
2836            prelude_adj[i] = 0.0;
2837            prelude_adj_dot[i] = 0.0;
2838            ror_dir_step(
2839                &self.prelude[i],
2840                i,
2841                prelude_vals,
2842                prelude_dot,
2843                prelude_adj,
2844                prelude_adj_dot,
2845                w,
2846                wd,
2847                1.0,
2848                out,
2849            );
2850        }
2851    }
2852
2853    /// Structural Hessian sparsity over the whole hybrid tape:
2854    /// every pair the prelude or any summand can produce.
2855    pub fn hessian_sparsity_all(&self) -> BTreeSet<(usize, usize)> {
2856        let mut pairs = hessian_sparsity_impl(&self.prelude);
2857
2858        // Per-prelude-slot var-set, reused across summands as the
2859        // var-set carrier for Shared refs.
2860        let prelude_var_sets = compute_var_sets(&self.prelude);
2861
2862        for s in &self.summands {
2863            summand_sparsity(&s.ops, &prelude_var_sets, &mut pairs);
2864        }
2865        pairs
2866    }
2867}
2868
2869/// Whether [`HybridTape::build_multi`] can build `exprs`, i.e. none of
2870/// them uses an opcode the hybrid (partial-separability) path rejects:
2871/// comparisons, AND/OR/NOT, if-then-else, min/max lists, or AMPL external
2872/// function calls. `build_into_summand` *panics* on those, so a caller
2873/// choosing between the hybrid path and the flat [`Tape`] path has to ask
2874/// first — there is nothing to catch.
2875///
2876/// Walks with an explicit stack (expression DAGs from `.nl` files can be
2877/// deeper than the default thread stack) and visits each shared CSE body
2878/// once, so the cost is linear in the DAG rather than in its unfolding.
2879pub fn hybrid_supported(exprs: &[Expr]) -> bool {
2880    let mut stack: Vec<&Expr> = exprs.iter().collect();
2881    let mut seen_cse: HashSet<*const Expr> = HashSet::new();
2882    while let Some(e) = stack.pop() {
2883        match e {
2884            Expr::Const(_) | Expr::Var(_) => {}
2885            Expr::Binary(_, a, b) => {
2886                stack.push(a);
2887                stack.push(b);
2888            }
2889            Expr::Unary(_, a) => stack.push(a),
2890            Expr::Sum(args) => stack.extend(args.iter()),
2891            Expr::Cse(body) => {
2892                if seen_cse.insert(Arc::as_ptr(body)) {
2893                    stack.push(body);
2894                }
2895            }
2896            Expr::Compare(..)
2897            | Expr::And(..)
2898            | Expr::Or(..)
2899            | Expr::Not(_)
2900            | Expr::Cond { .. }
2901            | Expr::MinList(_)
2902            | Expr::MaxList(_)
2903            | Expr::Funcall { .. } => return false,
2904        }
2905    }
2906    true
2907}
2908
2909/// True when `expr` (or any subexpression) is an AMPL external function
2910/// call. The hybrid summand path rejects funcalls outright, but the
2911/// *promoted*-CSE branch emits a shared CSE body via `build_recursive`
2912/// with an **empty** `ExternalResolver::default()` — it has no resolver
2913/// of its own. Without this pre-scan a funcall buried in a promoted CSE
2914/// would reach `build_recursive`'s `Expr::Funcall` arm and panic with the
2915/// misleading `unresolved AMPL funcall id <n>` message, instead of the
2916/// clear "not supported on the hybrid path" message the non-promoted
2917/// summand path raises. Pre-scanning makes both paths report the same
2918/// reason. (Funcalls are unsupported on the hybrid path regardless of
2919/// whether the id would resolve, so this never rejects a buildable tape.)
2920///
2921/// Walks with an explicit stack and visits each `Cse` body once, so the
2922/// cost is linear in the DAG rather than in its unfolding (gh#1066: a
2923/// chain `e = e + h*(u - e*e)` has `2^K` paths over `O(K)` nodes, and the
2924/// unmemoized walk never returned at K = 40). It also stops at a body
2925/// already in `built` — the prelude map — because that body was emitted
2926/// by `build_recursive` under the empty resolver, which would have
2927/// panicked on a funcall in it. Across one `build_multi` that bounds the
2928/// total scan by the DAG too, not by DAG times promoted bodies.
2929fn cse_contains_funcall(expr: &Expr, built: &HashMap<*const Expr, usize>) -> bool {
2930    let mut stack: Vec<&Expr> = vec![expr];
2931    let mut seen: HashSet<*const Expr> = HashSet::new();
2932    while let Some(e) = stack.pop() {
2933        match e {
2934            Expr::Funcall { .. } => return true,
2935            Expr::Const(_) | Expr::Var(_) => {}
2936            Expr::Binary(_, a, b) | Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2937                stack.push(a);
2938                stack.push(b);
2939            }
2940            Expr::Unary(_, a) | Expr::Not(a) => stack.push(a),
2941            Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2942                stack.extend(args.iter())
2943            }
2944            Expr::Cond { cond, then_, else_ } => {
2945                stack.push(cond);
2946                stack.push(then_);
2947                stack.push(else_);
2948            }
2949            Expr::Cse(body) => {
2950                let key = Arc::as_ptr(body);
2951                if !built.contains_key(&key) && seen.insert(key) {
2952                    stack.push(body);
2953                }
2954            }
2955        }
2956    }
2957    false
2958}
2959
2960/// Pass-1 helper: per-root walk that increments `counts[ptr]` the
2961/// first time a Cse pointer is encountered in this root. Recursing
2962/// into the body is gated on the first visit to avoid quadratic
2963/// blowup on heavily shared CSE DAGs.
2964fn count_cse_appearances(
2965    e: &Expr,
2966    seen_in_root: &mut HashSet<*const Expr>,
2967    counts: &mut HashMap<*const Expr, usize>,
2968) {
2969    match e {
2970        Expr::Const(_) | Expr::Var(_) => {}
2971        Expr::Binary(_, a, b) => {
2972            count_cse_appearances(a, seen_in_root, counts);
2973            count_cse_appearances(b, seen_in_root, counts);
2974        }
2975        Expr::Unary(_, a) => count_cse_appearances(a, seen_in_root, counts),
2976        Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2977            for a in args {
2978                count_cse_appearances(a, seen_in_root, counts);
2979            }
2980        }
2981        Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2982            count_cse_appearances(a, seen_in_root, counts);
2983            count_cse_appearances(b, seen_in_root, counts);
2984        }
2985        Expr::Not(a) => count_cse_appearances(a, seen_in_root, counts),
2986        Expr::Cond { cond, then_, else_ } => {
2987            count_cse_appearances(cond, seen_in_root, counts);
2988            count_cse_appearances(then_, seen_in_root, counts);
2989            count_cse_appearances(else_, seen_in_root, counts);
2990        }
2991        Expr::Cse(body) => {
2992            let key = Arc::as_ptr(body) as *const Expr;
2993            if seen_in_root.insert(key) {
2994                *counts.entry(key).or_insert(0) += 1;
2995                count_cse_appearances(body, seen_in_root, counts);
2996            }
2997        }
2998        Expr::Funcall { args, .. } => {
2999            for arg in args {
3000                if let FuncallArg::Real(e) = arg {
3001                    count_cse_appearances(e, seen_in_root, counts);
3002                }
3003            }
3004        }
3005    }
3006}
3007
3008/// Recursive summand builder. CSEs that meet the promotion bar
3009/// (≥ 2 roots reference them per `cse_count`) get a single prelude
3010/// emission via `build_recursive`; the summand records a Shared op
3011/// pointing at the prelude slot. Non-promoted CSEs are inlined
3012/// into the summand with intra-summand Arc-pointer dedup.
3013fn build_into_summand(
3014    expr: &Expr,
3015    local: &mut Vec<SummandOp>,
3016    local_cache: &mut HashMap<*const Expr, usize>,
3017    prelude: &mut Vec<TapeOp>,
3018    prelude_map: &mut HashMap<*const Expr, usize>,
3019    cse_count: &HashMap<*const Expr, usize>,
3020) -> usize {
3021    match expr {
3022        Expr::Const(c) => {
3023            let i = local.len();
3024            local.push(SummandOp::Local(TapeOp::Const(*c)));
3025            i
3026        }
3027        Expr::Var(j) => {
3028            let i = local.len();
3029            local.push(SummandOp::Local(TapeOp::Var(*j)));
3030            i
3031        }
3032        Expr::Binary(op, a, b) => {
3033            if let BinOp::Pow = op {
3034                if let Some(c) = peek_const(b) {
3035                    if let Some(i) = try_emit_const_pow_summand(
3036                        a,
3037                        c,
3038                        local,
3039                        local_cache,
3040                        prelude,
3041                        prelude_map,
3042                        cse_count,
3043                    ) {
3044                        return i;
3045                    }
3046                }
3047            }
3048            let l = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3049            let r = build_into_summand(b, local, local_cache, prelude, prelude_map, cse_count);
3050            let i = local.len();
3051            local.push(SummandOp::Local(match op {
3052                BinOp::Add => TapeOp::Add(l, r),
3053                BinOp::Sub => TapeOp::Sub(l, r),
3054                BinOp::Mul => TapeOp::Mul(l, r),
3055                BinOp::Div => TapeOp::Div(l, r),
3056                BinOp::Pow => TapeOp::Pow(l, r),
3057                BinOp::Atan2 => TapeOp::Atan2(l, r),
3058                BinOp::CEntropy => TapeOp::CEntropy(l, r),
3059            }));
3060            i
3061        }
3062        Expr::Unary(op, a) => {
3063            let v = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3064            let i = local.len();
3065            local.push(SummandOp::Local(match op {
3066                UnaryOp::Neg => TapeOp::Neg(v),
3067                UnaryOp::Sqrt => TapeOp::Sqrt(v),
3068                UnaryOp::Log => TapeOp::Log(v),
3069                UnaryOp::Log10 => TapeOp::Log10(v),
3070                UnaryOp::Exp => TapeOp::Exp(v),
3071                UnaryOp::Abs => TapeOp::Abs(v),
3072                UnaryOp::Sin => TapeOp::Sin(v),
3073                UnaryOp::Cos => TapeOp::Cos(v),
3074                UnaryOp::Tan => TapeOp::Tan(v),
3075                UnaryOp::Atan => TapeOp::Atan(v),
3076                UnaryOp::Acos => TapeOp::Acos(v),
3077                UnaryOp::Sinh => TapeOp::Sinh(v),
3078                UnaryOp::Cosh => TapeOp::Cosh(v),
3079                UnaryOp::Tanh => TapeOp::Tanh(v),
3080                UnaryOp::Asin => TapeOp::Asin(v),
3081                UnaryOp::Acosh => TapeOp::Acosh(v),
3082                UnaryOp::Asinh => TapeOp::Asinh(v),
3083                UnaryOp::Atanh => TapeOp::Atanh(v),
3084                UnaryOp::Erf => TapeOp::Erf(v),
3085                UnaryOp::XLogX => TapeOp::XLogX(v),
3086            }));
3087            i
3088        }
3089        Expr::Sum(args) => {
3090            if args.is_empty() {
3091                let i = local.len();
3092                local.push(SummandOp::Local(TapeOp::Const(0.0)));
3093                return i;
3094            }
3095            let mut acc = build_into_summand(
3096                &args[0],
3097                local,
3098                local_cache,
3099                prelude,
3100                prelude_map,
3101                cse_count,
3102            );
3103            for a in &args[1..] {
3104                let nxt =
3105                    build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3106                let i = local.len();
3107                local.push(SummandOp::Local(TapeOp::Add(acc, nxt)));
3108                acc = i;
3109            }
3110            acc
3111        }
3112        Expr::Cse(body) => {
3113            let key = Arc::as_ptr(body) as *const Expr;
3114            if let Some(&li) = local_cache.get(&key) {
3115                return li;
3116            }
3117            let promoted = cse_count.get(&key).copied().unwrap_or(0) >= 2;
3118            if promoted {
3119                // `build_recursive` below runs with an empty resolver, so a
3120                // funcall hidden inside the promoted body would panic with the
3121                // misleading "unresolved AMPL funcall id" message rather than
3122                // the clear hybrid-unsupported message the non-promoted summand
3123                // path (and the `Expr::Funcall` arm at the bottom) raises.
3124                // Reject it up front so both CSE paths report the same reason.
3125                if cse_contains_funcall(body, prelude_map) {
3126                    panic!(
3127                        "HybridTape: AMPL external function calls are not supported on the \
3128                         hybrid (partial-separability) tape path. Build with \
3129                         Tape::build_with_externals instead."
3130                    );
3131                }
3132                // Build (or reuse) the prelude slot for this CSE.
3133                // `build_recursive(expr, ...)` hits the Cse arm,
3134                // emits the body once into prelude, and caches it
3135                // in `prelude_map` keyed by this Arc pointer.
3136                let pslot =
3137                    build_recursive(expr, prelude, prelude_map, &ExternalResolver::default());
3138                let li = local.len();
3139                local.push(SummandOp::Shared(pslot));
3140                local_cache.insert(key, li);
3141                li
3142            } else {
3143                let li =
3144                    build_into_summand(body, local, local_cache, prelude, prelude_map, cse_count);
3145                local_cache.insert(key, li);
3146                li
3147            }
3148        }
3149        Expr::Compare(_, _, _)
3150        | Expr::And(_, _)
3151        | Expr::Or(_, _)
3152        | Expr::Not(_)
3153        | Expr::Cond { .. }
3154        | Expr::MinList(_)
3155        | Expr::MaxList(_) => {
3156            panic!(
3157                "HybridTape: conditional / logical / min-max opcodes (comparisons, \
3158                 AND/OR/NOT, if-then-else, min/max lists) are not supported on the \
3159                 hybrid (partial-separability) tape path. Build with \
3160                 Tape::build_with_externals instead."
3161            );
3162        }
3163        Expr::Funcall { .. } => {
3164            panic!(
3165                "HybridTape: AMPL external function calls are not supported on the \
3166                 hybrid (partial-separability) tape path. Build with Tape::build_with_externals \
3167                 instead."
3168            );
3169        }
3170    }
3171}
3172
3173/// Pow-lowering specialised for summand builds. Mirrors
3174/// `try_emit_const_pow` but with summand-flavoured emission.
3175fn try_emit_const_pow_summand(
3176    base_expr: &Expr,
3177    c: f64,
3178    local: &mut Vec<SummandOp>,
3179    local_cache: &mut HashMap<*const Expr, usize>,
3180    prelude: &mut Vec<TapeOp>,
3181    prelude_map: &mut HashMap<*const Expr, usize>,
3182    cse_count: &HashMap<*const Expr, usize>,
3183) -> Option<usize> {
3184    if c == 0.0 {
3185        let i = local.len();
3186        local.push(SummandOp::Local(TapeOp::Const(1.0)));
3187        return Some(i);
3188    }
3189    if c == 1.0 {
3190        return Some(build_into_summand(
3191            base_expr,
3192            local,
3193            local_cache,
3194            prelude,
3195            prelude_map,
3196            cse_count,
3197        ));
3198    }
3199    if c == 0.5 {
3200        let b = build_into_summand(
3201            base_expr,
3202            local,
3203            local_cache,
3204            prelude,
3205            prelude_map,
3206            cse_count,
3207        );
3208        let i = local.len();
3209        local.push(SummandOp::Local(TapeOp::Sqrt(b)));
3210        return Some(i);
3211    }
3212    if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
3213        let n = c.abs() as u32;
3214        if n == 0 {
3215            let i = local.len();
3216            local.push(SummandOp::Local(TapeOp::Const(1.0)));
3217            return Some(i);
3218        }
3219        let b = build_into_summand(
3220            base_expr,
3221            local,
3222            local_cache,
3223            prelude,
3224            prelude_map,
3225            cse_count,
3226        );
3227        let pos = emit_int_pow_summand(b, n, local);
3228        if c < 0.0 {
3229            let one_idx = local.len();
3230            local.push(SummandOp::Local(TapeOp::Const(1.0)));
3231            let i = local.len();
3232            local.push(SummandOp::Local(TapeOp::Div(one_idx, pos)));
3233            return Some(i);
3234        }
3235        return Some(pos);
3236    }
3237    None
3238}
3239
3240fn emit_int_pow_summand(base: usize, n: u32, local: &mut Vec<SummandOp>) -> usize {
3241    debug_assert!(n >= 1);
3242    if n == 1 {
3243        return base;
3244    }
3245    let half = emit_int_pow_summand(base, n / 2, local);
3246    let squared = local.len();
3247    local.push(SummandOp::Local(TapeOp::Mul(half, half)));
3248    if n % 2 == 1 {
3249        let i = local.len();
3250        local.push(SummandOp::Local(TapeOp::Mul(squared, base)));
3251        i
3252    } else {
3253        squared
3254    }
3255}
3256
3257/// Walk a summand's local op DAG from `root`, returning the
3258/// reachable local slots (sorted ascending) plus the distinct
3259/// prelude slots referenced by any Shared op along the way.
3260fn compute_local_reach(ops: &[SummandOp], root: usize) -> (Vec<usize>, Vec<usize>) {
3261    let mut visited = vec![false; ops.len()];
3262    let mut reach: Vec<usize> = Vec::new();
3263    let mut shared: BTreeSet<usize> = BTreeSet::new();
3264    let mut stack: Vec<usize> = Vec::with_capacity(16);
3265    visited[root] = true;
3266    reach.push(root);
3267    stack.push(root);
3268    while let Some(s) = stack.pop() {
3269        match &ops[s] {
3270            SummandOp::Local(op) => {
3271                let (a, b) = op_operands(op);
3272                if let Some(a) = a {
3273                    if !visited[a] {
3274                        visited[a] = true;
3275                        reach.push(a);
3276                        stack.push(a);
3277                    }
3278                }
3279                if let Some(b) = b {
3280                    if !visited[b] {
3281                        visited[b] = true;
3282                        reach.push(b);
3283                        stack.push(b);
3284                    }
3285                }
3286            }
3287            SummandOp::Shared(k) => {
3288                shared.insert(*k);
3289            }
3290        }
3291    }
3292    reach.sort_unstable();
3293    (reach, shared.into_iter().collect())
3294}
3295
3296/// Epoch-tagged BFS over the prelude operand DAG, accumulating
3297/// reachable slots into `out`. Caller is responsible for sorting
3298/// `out` after a batch of starts has been processed.
3299fn bfs_prelude(
3300    prelude: &[TapeOp],
3301    start: usize,
3302    visited: &mut [u32],
3303    cur: u32,
3304    stack: &mut Vec<usize>,
3305    out: &mut Vec<usize>,
3306) {
3307    if visited[start] == cur {
3308        return;
3309    }
3310    visited[start] = cur;
3311    out.push(start);
3312    stack.push(start);
3313    while let Some(s) = stack.pop() {
3314        let (a, b) = op_operands(&prelude[s]);
3315        if let Some(a) = a {
3316            if visited[a] != cur {
3317                visited[a] = cur;
3318                out.push(a);
3319                stack.push(a);
3320            }
3321        }
3322        if let Some(b) = b {
3323            if visited[b] != cur {
3324                visited[b] = cur;
3325                out.push(b);
3326                stack.push(b);
3327            }
3328        }
3329    }
3330}
3331
3332/// Per-op var-set for the prelude — every slot's transitive
3333/// variable footprint. Used by `summand_sparsity` to expand
3334/// `SummandOp::Shared(k)` into its var-set carrier.
3335fn compute_var_sets(ops: &[TapeOp]) -> Vec<BTreeSet<usize>> {
3336    let mut out: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3337    for op in ops {
3338        let vs: BTreeSet<usize> = match op {
3339            TapeOp::Const(_) => BTreeSet::new(),
3340            TapeOp::Var(j) => {
3341                let mut s = BTreeSet::new();
3342                s.insert(*j);
3343                s
3344            }
3345            TapeOp::Add(a, b)
3346            | TapeOp::Sub(a, b)
3347            | TapeOp::Mul(a, b)
3348            | TapeOp::Div(a, b)
3349            | TapeOp::Pow(a, b)
3350            | TapeOp::Atan2(a, b)
3351            | TapeOp::CEntropy(a, b) => out[*a].union(&out[*b]).copied().collect(),
3352            TapeOp::Neg(a)
3353            | TapeOp::Abs(a)
3354            | TapeOp::Sqrt(a)
3355            | TapeOp::Exp(a)
3356            | TapeOp::Log(a)
3357            | TapeOp::Log10(a)
3358            | TapeOp::Sin(a)
3359            | TapeOp::Cos(a)
3360            | TapeOp::Tan(a)
3361            | TapeOp::Atan(a)
3362            | TapeOp::Acos(a)
3363            | TapeOp::Sinh(a)
3364            | TapeOp::Cosh(a)
3365            | TapeOp::Tanh(a)
3366            | TapeOp::Asin(a)
3367            | TapeOp::Acosh(a)
3368            | TapeOp::Asinh(a)
3369            | TapeOp::Erf(a)
3370            | TapeOp::XLogX(a)
3371            | TapeOp::Atanh(a) => out[*a].clone(),
3372            TapeOp::Cmp(_, _, _)
3373            | TapeOp::And(_, _)
3374            | TapeOp::Or(_, _)
3375            | TapeOp::Not(_)
3376            | TapeOp::Select(_, _, _)
3377            | TapeOp::Min(_, _)
3378            | TapeOp::Max(_, _) => unreachable!(
3379                "HybridTape prelude cannot contain conditional / logical / min-max \
3380                 TapeOps; build_into_summand panics on those Expr variants."
3381            ),
3382            TapeOp::Funcall(_) => unreachable!(
3383                "HybridTape prelude cannot contain TapeOp::Funcall; \
3384                 build_into_summand panics on Expr::Funcall."
3385            ),
3386        };
3387        out.push(vs);
3388    }
3389    out
3390}
3391
3392/// Per-op Hessian-sparsity propagation over a summand's mixed
3393/// SummandOp slice. Shared refs contribute their prelude var-set
3394/// but do not themselves emit pairs (those came from
3395/// `hessian_sparsity_impl(&prelude)`).
3396fn summand_sparsity(
3397    ops: &[SummandOp],
3398    prelude_var_sets: &[BTreeSet<usize>],
3399    pairs: &mut BTreeSet<(usize, usize)>,
3400) {
3401    let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3402    let emit_cross =
3403        |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3404            for &v1 in s1 {
3405                for &v2 in s2 {
3406                    let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
3407                    pairs.insert((r, c));
3408                }
3409            }
3410        };
3411    let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3412        let vars: Vec<usize> = s.iter().copied().collect();
3413        for (ai, &vi) in vars.iter().enumerate() {
3414            for &vj in &vars[..=ai] {
3415                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
3416                pairs.insert((r, c));
3417            }
3418        }
3419    };
3420    for so in ops {
3421        let vset: BTreeSet<usize> = match so {
3422            SummandOp::Shared(k) => prelude_var_sets[*k].clone(),
3423            SummandOp::Local(op) => match op {
3424                TapeOp::Const(_) => BTreeSet::new(),
3425                TapeOp::Var(j) => {
3426                    let mut s = BTreeSet::new();
3427                    s.insert(*j);
3428                    s
3429                }
3430                TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
3431                    var_sets[*a].union(&var_sets[*b]).copied().collect()
3432                }
3433                TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
3434                TapeOp::Mul(a, b) => {
3435                    emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3436                    var_sets[*a].union(&var_sets[*b]).copied().collect()
3437                }
3438                TapeOp::Div(a, b) => {
3439                    emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3440                    emit_self(&var_sets[*b], pairs);
3441                    var_sets[*a].union(&var_sets[*b]).copied().collect()
3442                }
3443                TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
3444                    let combined: BTreeSet<usize> =
3445                        var_sets[*a].union(&var_sets[*b]).copied().collect();
3446                    emit_self(&combined, pairs);
3447                    combined
3448                }
3449                TapeOp::Sqrt(a)
3450                | TapeOp::Exp(a)
3451                | TapeOp::Log(a)
3452                | TapeOp::Log10(a)
3453                | TapeOp::Sin(a)
3454                | TapeOp::Cos(a)
3455                | TapeOp::Tan(a)
3456                | TapeOp::Atan(a)
3457                | TapeOp::Acos(a)
3458                | TapeOp::Sinh(a)
3459                | TapeOp::Cosh(a)
3460                | TapeOp::Tanh(a)
3461                | TapeOp::Asin(a)
3462                | TapeOp::Acosh(a)
3463                | TapeOp::Asinh(a)
3464                | TapeOp::Erf(a)
3465                | TapeOp::XLogX(a)
3466                | TapeOp::Atanh(a) => {
3467                    emit_self(&var_sets[*a], pairs);
3468                    var_sets[*a].clone()
3469                }
3470                TapeOp::Cmp(_, _, _)
3471                | TapeOp::And(_, _)
3472                | TapeOp::Or(_, _)
3473                | TapeOp::Not(_)
3474                | TapeOp::Select(_, _, _)
3475                | TapeOp::Min(_, _)
3476                | TapeOp::Max(_, _) => unreachable!(
3477                    "HybridTape summand cannot contain conditional / logical / min-max \
3478                     TapeOps; build_into_summand panics on those Expr variants."
3479                ),
3480                TapeOp::Funcall(_) => unreachable!(
3481                    "HybridTape summand cannot contain TapeOp::Funcall; \
3482                     build_into_summand panics on Expr::Funcall."
3483                ),
3484            },
3485        };
3486        var_sets.push(vset);
3487    }
3488}
3489
3490/// Operand indices of a `TapeOp`, normalized into a fixed-length
3491/// array so callers don't need to re-match every site.
3492#[inline]
3493pub(crate) fn op_operands(op: &TapeOp) -> (Option<usize>, Option<usize>) {
3494    match op {
3495        TapeOp::Const(_) | TapeOp::Var(_) => (None, None),
3496        TapeOp::Add(a, b)
3497        | TapeOp::Sub(a, b)
3498        | TapeOp::Mul(a, b)
3499        | TapeOp::Div(a, b)
3500        | TapeOp::Pow(a, b)
3501        | TapeOp::Atan2(a, b)
3502        | TapeOp::CEntropy(a, b) => (Some(*a), Some(*b)),
3503        TapeOp::Neg(a)
3504        | TapeOp::Abs(a)
3505        | TapeOp::Sqrt(a)
3506        | TapeOp::Exp(a)
3507        | TapeOp::Log(a)
3508        | TapeOp::Log10(a)
3509        | TapeOp::Sin(a)
3510        | TapeOp::Cos(a)
3511        | TapeOp::Tan(a)
3512        | TapeOp::Atan(a)
3513        | TapeOp::Acos(a)
3514        | TapeOp::Sinh(a)
3515        | TapeOp::Cosh(a)
3516        | TapeOp::Tanh(a)
3517        | TapeOp::Asin(a)
3518        | TapeOp::Acosh(a)
3519        | TapeOp::Asinh(a)
3520        | TapeOp::Erf(a)
3521        | TapeOp::XLogX(a)
3522        | TapeOp::Atanh(a) => (Some(*a), None),
3523        // Conditional / logical TapeOps never reach the HybridTape
3524        // operand-walk (build_into_summand rejects them). Cmp/And/Or
3525        // have two operands; Not has one; Select's three can't be
3526        // expressed in this two-slot shape, so it would be a bug to
3527        // see it here.
3528        TapeOp::Cmp(_, a, b) | TapeOp::And(a, b) | TapeOp::Or(a, b) => (Some(*a), Some(*b)),
3529        TapeOp::Not(a) => (Some(*a), None),
3530        TapeOp::Select(_, _, _) => unreachable!(
3531            "op_operands: TapeOp::Select has three operands and is unsupported on \
3532             the HybridTape path"
3533        ),
3534        TapeOp::Min(_, _) | TapeOp::Max(_, _) => unreachable!(
3535            "op_operands: TapeOp::Min/Max are unsupported on the HybridTape path \
3536             (build_into_summand rejects min/max lists)"
3537        ),
3538        // Returning `(None, None)` here would be a silent wrong answer, not a
3539        // conservative one: a Funcall's tape arguments would be missing from
3540        // `local_reach`, and the directional Hessian deliberately does NOT
3541        // pre-zero `local_dot`, so those slots would be read as stale scratch
3542        // from a previous summand. Unreachable today —
3543        // `build_into_summand` panics on `Expr::Funcall` before any tape is
3544        // built — and this keeps it loud if that ever changes.
3545        TapeOp::Funcall(_) => unreachable!(
3546            "op_operands: TapeOp::Funcall is unsupported on the HybridTape path \
3547             (build_into_summand rejects external function calls)"
3548        ),
3549    }
3550}
3551
3552fn vars_in(ops: &[TapeOp], reach: &[usize]) -> Vec<usize> {
3553    let mut s: BTreeSet<usize> = BTreeSet::new();
3554    for &i in reach {
3555        if let TapeOp::Var(j) = &ops[i] {
3556            s.insert(*j);
3557        }
3558    }
3559    s.into_iter().collect()
3560}
3561
3562// ----- Free-function AD step kernels used by GlobalTape -----
3563
3564#[inline]
3565fn fwd_step(op: &TapeOp, x: &[f64], vals: &[f64]) -> f64 {
3566    match op {
3567        TapeOp::Const(c) => *c,
3568        TapeOp::Var(i) => x[*i],
3569        TapeOp::Add(a, b) => vals[*a] + vals[*b],
3570        TapeOp::Sub(a, b) => vals[*a] - vals[*b],
3571        TapeOp::Mul(a, b) => vals[*a] * vals[*b],
3572        TapeOp::Div(a, b) => vals[*a] / vals[*b],
3573        TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
3574        TapeOp::Neg(a) => -vals[*a],
3575        TapeOp::Abs(a) => vals[*a].abs(),
3576        TapeOp::Sqrt(a) => vals[*a].sqrt(),
3577        TapeOp::Exp(a) => vals[*a].exp(),
3578        TapeOp::Log(a) => vals[*a].ln(),
3579        TapeOp::Log10(a) => vals[*a].log10(),
3580        TapeOp::Sin(a) => vals[*a].sin(),
3581        TapeOp::Cos(a) => vals[*a].cos(),
3582        TapeOp::Tan(a) => vals[*a].tan(),
3583        TapeOp::Atan(a) => vals[*a].atan(),
3584        TapeOp::Acos(a) => vals[*a].acos(),
3585        TapeOp::Sinh(a) => vals[*a].sinh(),
3586        TapeOp::Cosh(a) => vals[*a].cosh(),
3587        TapeOp::Tanh(a) => vals[*a].tanh(),
3588        TapeOp::Asin(a) => vals[*a].asin(),
3589        TapeOp::Acosh(a) => vals[*a].acosh(),
3590        TapeOp::Asinh(a) => vals[*a].asinh(),
3591        TapeOp::Atanh(a) => atanh(vals[*a]),
3592        TapeOp::Erf(a) => erf(vals[*a]),
3593        TapeOp::XLogX(a) => xlogx(vals[*a]),
3594        TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
3595        TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
3596        TapeOp::Cmp(_, _, _)
3597        | TapeOp::And(_, _)
3598        | TapeOp::Or(_, _)
3599        | TapeOp::Not(_)
3600        | TapeOp::Select(_, _, _)
3601        | TapeOp::Min(_, _)
3602        | TapeOp::Max(_, _) => panic!(
3603            "GlobalTape free-function kernels do not implement conditional / logical \
3604             / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3605             instead."
3606        ),
3607        TapeOp::Funcall(fc) => {
3608            let FuncallData { lib, name, args } = fc.as_ref();
3609            let call_args = funcall_to_ext_args(args, vals);
3610            let res = lib
3611                .eval(name, &call_args, false, false)
3612                .unwrap_or_else(|e| panic!("external function '{name}' eval failed: {e}"));
3613            res.value
3614        }
3615    }
3616}
3617
3618#[inline]
3619fn rev_step(op: &TapeOp, i: usize, vals: &[f64], adj: &mut [f64], a: f64, grad: &mut [f64]) {
3620    match op {
3621        TapeOp::Const(_) => {}
3622        TapeOp::Var(j) => {
3623            grad[*j] += a;
3624        }
3625        TapeOp::Add(l, r) => {
3626            adj[*l] += a;
3627            adj[*r] += a;
3628        }
3629        TapeOp::Sub(l, r) => {
3630            adj[*l] += a;
3631            adj[*r] -= a;
3632        }
3633        TapeOp::Mul(l, r) => {
3634            adj[*l] += a * vals[*r];
3635            adj[*r] += a * vals[*l];
3636        }
3637        TapeOp::Div(l, r) => {
3638            // Kahan form -- see the identical arm in the reverse sweep above.
3639            let rv = vals[*r];
3640            adj[*l] += a / rv;
3641            adj[*r] -= a * vals[i] / rv;
3642        }
3643        TapeOp::Pow(l, r) => {
3644            let lv = vals[*l];
3645            let rv = vals[*r];
3646            if rv != 0.0 {
3647                adj[*l] += a * rv * lv.powf(rv - 1.0);
3648            }
3649            if lv > 0.0 {
3650                adj[*r] += a * vals[i] * lv.ln();
3651            }
3652        }
3653        TapeOp::Neg(j) => {
3654            adj[*j] -= a;
3655        }
3656        TapeOp::Abs(j) => {
3657            if vals[*j] >= 0.0 {
3658                adj[*j] += a;
3659            } else {
3660                adj[*j] -= a;
3661            }
3662        }
3663        TapeOp::Sqrt(j) => {
3664            let sv = vals[i];
3665            if sv > 0.0 {
3666                adj[*j] += a * 0.5 / sv;
3667            }
3668        }
3669        TapeOp::Exp(j) => {
3670            adj[*j] += a * vals[i];
3671        }
3672        TapeOp::Log(j) => {
3673            adj[*j] += a / vals[*j];
3674        }
3675        TapeOp::Log10(j) => {
3676            adj[*j] += div_u_ln10(a, vals[*j]);
3677        }
3678        TapeOp::Sin(j) => {
3679            adj[*j] += a * vals[*j].cos();
3680        }
3681        TapeOp::Cos(j) => {
3682            adj[*j] -= a * vals[*j].sin();
3683        }
3684        TapeOp::Tan(j) => {
3685            let t = vals[i];
3686            adj[*j] += a * (1.0 + t * t);
3687        }
3688        TapeOp::Atan(j) => {
3689            adj[*j] += a * atan_d1(vals[*j]);
3690        }
3691        TapeOp::Acos(j) => {
3692            adj[*j] += a * acos_d1(vals[*j]);
3693        }
3694        TapeOp::Sinh(j) => {
3695            adj[*j] += a * vals[*j].cosh();
3696        }
3697        TapeOp::Cosh(j) => {
3698            adj[*j] += a * vals[*j].sinh();
3699        }
3700        TapeOp::Tanh(j) => {
3701            adj[*j] += a * tanh_d1(vals[*j]);
3702        }
3703        TapeOp::Asin(j) => {
3704            adj[*j] += a * asin_d1(vals[*j]);
3705        }
3706        TapeOp::Acosh(j) => {
3707            adj[*j] += a * acosh_d1(vals[*j]);
3708        }
3709        TapeOp::Asinh(j) => {
3710            adj[*j] += a * asinh_d1(vals[*j]);
3711        }
3712        TapeOp::Atanh(j) => {
3713            adj[*j] += a * atanh_d1(vals[*j]);
3714        }
3715        TapeOp::Erf(j) => {
3716            adj[*j] += a * erf_d1(vals[*j]);
3717        }
3718        TapeOp::XLogX(j) => {
3719            adj[*j] += a * xlogx_d1(vals[*j]);
3720        }
3721        TapeOp::CEntropy(l, r) => {
3722            adj[*l] += a * centropy_da(vals[*l], vals[*r]);
3723            adj[*r] += a * centropy_db(vals[*l], vals[*r]);
3724        }
3725        TapeOp::Atan2(l, r) => {
3726            let y = vals[*l];
3727            let x = vals[*r];
3728            let (fa, fb) = atan2_d1(y, x);
3729            adj[*l] += a * fa;
3730            adj[*r] += a * fb;
3731        }
3732        TapeOp::Cmp(_, _, _)
3733        | TapeOp::And(_, _)
3734        | TapeOp::Or(_, _)
3735        | TapeOp::Not(_)
3736        | TapeOp::Select(_, _, _)
3737        | TapeOp::Min(_, _)
3738        | TapeOp::Max(_, _) => panic!(
3739            "GlobalTape free-function kernels do not implement conditional / logical \
3740             / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3741             instead."
3742        ),
3743        TapeOp::Funcall(fc) => {
3744            let FuncallData { lib, name, args } = fc.as_ref();
3745            let call_args = funcall_to_ext_args(args, vals);
3746            let res = lib
3747                .eval(name, &call_args, true, false)
3748                .unwrap_or_else(|e| panic!("external function '{name}' reverse eval failed: {e}"));
3749            let derivs = res.derivs.expect("want_derivs=true returns derivs");
3750            let mut k = 0usize;
3751            for arg in args {
3752                if let TapeFuncallArg::Tape(idx) = arg {
3753                    adj[*idx] += a * derivs[k];
3754                    k += 1;
3755                }
3756            }
3757            let _ = i;
3758            let _ = grad;
3759        }
3760    }
3761}
3762
3763/// Directional forward-tangent step: [`fwd_tan_step`] with a dense
3764/// seed vector instead of a single seed variable. The per-op chain
3765/// rule is identical (and identical to the inline arms of
3766/// [`Tape::hessian_directional`], which the hybrid Hessian path must
3767/// agree with); only the `Var` arm differs — it reads `seed[k]`
3768/// rather than testing `k == seed_var` — which is what lets one pass
3769/// carry a whole Hessian color.
3770#[inline]
3771fn fwd_dir_step(op: &TapeOp, seed: &[f64], vals: &[f64], dot: &[f64], i: usize) -> f64 {
3772    match op {
3773        TapeOp::Const(_) => 0.0,
3774        TapeOp::Var(k) => seed[*k],
3775        TapeOp::Add(a, b) => dot[*a] + dot[*b],
3776        TapeOp::Sub(a, b) => dot[*a] - dot[*b],
3777        TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
3778        TapeOp::Div(a, b) => {
3779            // Kahan form -- see the identical arm in `forward_tangent` for why the
3780            // squared denominator loses representable tangents at extreme |b|.
3781            (dot[*a] - vals[i] * dot[*b]) / vals[*b]
3782        }
3783        TapeOp::Pow(a, b) => {
3784            let u = vals[*a];
3785            let r = vals[*b];
3786            let du = dot[*a];
3787            let dr = dot[*b];
3788            let mut result = 0.0;
3789            // Match the reverse-mode gradient's guard (`rv != 0.0` only): at base
3790            // u == 0 the slope is still well defined for r >= 1 (and a genuine
3791            // ±inf for r < 1), so it must not be silently dropped, or the forward
3792            // tangent disagrees with the reverse gradient.
3793            if r != 0.0 {
3794                result += r * u.powf(r - 1.0) * du;
3795            }
3796            if u > 0.0 {
3797                result += vals[i] * u.ln() * dr;
3798            }
3799            result
3800        }
3801        TapeOp::Neg(a) => -dot[*a],
3802        TapeOp::Abs(a) => {
3803            if vals[*a] >= 0.0 {
3804                dot[*a]
3805            } else {
3806                -dot[*a]
3807            }
3808        }
3809        TapeOp::Sqrt(a) => {
3810            let sv = vals[i];
3811            if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
3812        }
3813        TapeOp::Exp(a) => vals[i] * dot[*a],
3814        TapeOp::Log(a) => dot[*a] / vals[*a],
3815        TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
3816        TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
3817        TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
3818        TapeOp::Tan(a) => {
3819            let t = vals[i];
3820            (1.0 + t * t) * dot[*a]
3821        }
3822        TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
3823        TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
3824        TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
3825        TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
3826        TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
3827        TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
3828        TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
3829        TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
3830        TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
3831        TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
3832        TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
3833        TapeOp::CEntropy(a, b) => {
3834            centropy_da(vals[*a], vals[*b]) * dot[*a] + centropy_db(vals[*a], vals[*b]) * dot[*b]
3835        }
3836        TapeOp::Atan2(a, b) => {
3837            let y = vals[*a];
3838            let x = vals[*b];
3839            let (fa, fb) = atan2_d1(y, x);
3840            fa * dot[*a] + fb * dot[*b]
3841        }
3842        TapeOp::Cmp(_, _, _)
3843        | TapeOp::And(_, _)
3844        | TapeOp::Or(_, _)
3845        | TapeOp::Not(_)
3846        | TapeOp::Select(_, _, _)
3847        | TapeOp::Min(_, _)
3848        | TapeOp::Max(_, _) => panic!(
3849            "GlobalTape free-function kernels do not implement conditional / logical \
3850             / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3851             instead."
3852        ),
3853        TapeOp::Funcall(fc) => {
3854            let FuncallData { lib, name, args } = fc.as_ref();
3855            let call_args = funcall_to_ext_args(args, vals);
3856            let res = lib
3857                .eval(name, &call_args, true, false)
3858                .unwrap_or_else(|e| panic!("external function '{name}' tangent eval failed: {e}"));
3859            let derivs = res.derivs.expect("want_derivs=true returns derivs");
3860            let mut acc = 0.0;
3861            let mut k = 0usize;
3862            for arg in args {
3863                if let TapeFuncallArg::Tape(idx) = arg {
3864                    acc += derivs[k] * dot[*idx];
3865                    k += 1;
3866                }
3867            }
3868            let _ = seed;
3869            acc
3870        }
3871    }
3872}
3873
3874/// Directional reverse-over-tangent step: [`ror_step`] with the
3875/// hash-map scatter replaced by the dense accumulation
3876/// [`Tape::hessian_directional`] uses — at `Var(k)` the emitted
3877/// second-order contribution is `out[k] += weight * wd`, landing in
3878/// the coloring's per-color `compressed` buffer with no hashing. All
3879/// other arms are the exact arithmetic of `hessian_directional`'s
3880/// reverse sweep, so the local part of the hybrid Hessian is
3881/// bit-identical to the flat tape's.
3882#[allow(clippy::too_many_arguments)]
3883#[inline]
3884fn ror_dir_step(
3885    op: &TapeOp,
3886    i: usize,
3887    vals: &[f64],
3888    dot: &[f64],
3889    adj: &mut [f64],
3890    adj_dot: &mut [f64],
3891    w: f64,
3892    wd: f64,
3893    weight: f64,
3894    out: &mut [f64],
3895) {
3896    match op {
3897        TapeOp::Const(_) => {}
3898        TapeOp::Var(k) => {
3899            if wd != 0.0 {
3900                out[*k] += weight * wd;
3901            }
3902        }
3903        TapeOp::Add(a, b) => {
3904            adj[*a] += w;
3905            adj[*b] += w;
3906            adj_dot[*a] += wd;
3907            adj_dot[*b] += wd;
3908        }
3909        TapeOp::Sub(a, b) => {
3910            adj[*a] += w;
3911            adj[*b] -= w;
3912            adj_dot[*a] += wd;
3913            adj_dot[*b] -= wd;
3914        }
3915        TapeOp::Mul(a, b) => {
3916            adj[*a] += w * vals[*b];
3917            adj[*b] += w * vals[*a];
3918            adj_dot[*a] += wd * vals[*b] + w * dot[*b];
3919            adj_dot[*b] += wd * vals[*a] + w * dot[*a];
3920        }
3921        TapeOp::Div(a, b) => {
3922            // Kahan form -- see the identical arm in the dense
3923            // reverse-over-forward sweep for the derivation.
3924            let vb = vals[*b];
3925            let q = vals[i];
3926            let qd = dot[i];
3927            adj[*a] += w / vb;
3928            adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
3929            adj[*b] -= w * q / vb;
3930            adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
3931        }
3932        TapeOp::Pow(a, b) => {
3933            let u = vals[*a];
3934            let r = vals[*b];
3935            let du = dot[*a];
3936            let dr = dot[*b];
3937            if r != 0.0 {
3938                if u != 0.0 {
3939                    let p_a = r * u.powf(r - 1.0);
3940                    adj[*a] += w * p_a;
3941                    let mut dp_a = dr * u.powf(r - 1.0);
3942                    if u > 0.0 {
3943                        dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
3944                    } else {
3945                        dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
3946                    }
3947                    adj_dot[*a] += wd * p_a + w * dp_a;
3948                } else if r >= 2.0 {
3949                    let p_a = 0.0;
3950                    adj[*a] += w * p_a;
3951                    let dp_a = if r == 2.0 {
3952                        2.0 * du
3953                    } else {
3954                        r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
3955                    };
3956                    adj_dot[*a] += wd * p_a + w * dp_a;
3957                }
3958            }
3959            if u > 0.0 {
3960                let ln_u = u.ln();
3961                let p_b = vals[i] * ln_u;
3962                adj[*b] += w * p_b;
3963                let dur = vals[i] * (r * du / u + dr * ln_u);
3964                let dp_b = dur * ln_u + vals[i] * du / u;
3965                adj_dot[*b] += wd * p_b + w * dp_b;
3966            }
3967        }
3968        TapeOp::Neg(a) => {
3969            adj[*a] -= w;
3970            adj_dot[*a] -= wd;
3971        }
3972        TapeOp::Abs(a) => {
3973            let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
3974            adj[*a] += w * s;
3975            adj_dot[*a] += wd * s;
3976        }
3977        TapeOp::Sqrt(a) => {
3978            let sv = vals[i];
3979            if sv > 0.0 {
3980                let fp = 0.5 / sv;
3981                let fpp = sqrt_d2(vals[*a], sv);
3982                adj[*a] += w * fp;
3983                adj_dot[*a] += wd * fp + w * fpp * dot[*a];
3984            }
3985        }
3986        TapeOp::Exp(a) => {
3987            let ev = vals[i];
3988            adj[*a] += w * ev;
3989            adj_dot[*a] += wd * ev + w * ev * dot[*a];
3990        }
3991        TapeOp::Log(a) => {
3992            let u = vals[*a];
3993            adj[*a] += w / u;
3994            adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
3995        }
3996        TapeOp::Log10(a) => {
3997            let u = vals[*a];
3998            adj[*a] += div_u_ln10(w, u);
3999            adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
4000        }
4001        TapeOp::Sin(a) => {
4002            let u = vals[*a];
4003            let cu = u.cos();
4004            adj[*a] += w * cu;
4005            adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
4006        }
4007        TapeOp::Cos(a) => {
4008            let u = vals[*a];
4009            let su = u.sin();
4010            adj[*a] -= w * su;
4011            adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
4012        }
4013        TapeOp::Tan(a) => {
4014            let t = vals[i];
4015            let gp = 1.0 + t * t;
4016            let gpp = 2.0 * t * gp;
4017            adj[*a] += w * gp;
4018            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4019        }
4020        TapeOp::Atan(a) => {
4021            let u = vals[*a];
4022            let gp = atan_d1(u);
4023            let gpp = atan_d2(u);
4024            adj[*a] += w * gp;
4025            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4026        }
4027        TapeOp::Acos(a) => {
4028            let u = vals[*a];
4029            let gp = acos_d1(u);
4030            let gpp = acos_d2(u);
4031            adj[*a] += w * gp;
4032            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4033        }
4034        TapeOp::Sinh(a) => {
4035            let u = vals[*a];
4036            let gp = u.cosh();
4037            let gpp = vals[i]; // sinh(u)
4038            adj[*a] += w * gp;
4039            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4040        }
4041        TapeOp::Cosh(a) => {
4042            let u = vals[*a];
4043            let gp = u.sinh();
4044            let gpp = vals[i]; // cosh(u)
4045            adj[*a] += w * gp;
4046            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4047        }
4048        TapeOp::Tanh(a) => {
4049            let t = vals[i];
4050            let gp = tanh_d1(vals[*a]);
4051            let gpp = -2.0 * t * gp;
4052            adj[*a] += w * gp;
4053            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4054        }
4055        TapeOp::Asin(a) => {
4056            let u = vals[*a];
4057            let gp = asin_d1(u);
4058            let gpp = asin_d2(u);
4059            adj[*a] += w * gp;
4060            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4061        }
4062        TapeOp::Acosh(a) => {
4063            let u = vals[*a];
4064            let gp = acosh_d1(u);
4065            let gpp = acosh_d2(u);
4066            adj[*a] += w * gp;
4067            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4068        }
4069        TapeOp::Asinh(a) => {
4070            let u = vals[*a];
4071            let gp = asinh_d1(u);
4072            let gpp = asinh_d2(u);
4073            adj[*a] += w * gp;
4074            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4075        }
4076        TapeOp::Atanh(a) => {
4077            let u = vals[*a];
4078            let gp = atanh_d1(u);
4079            let gpp = atanh_d2(u);
4080            adj[*a] += w * gp;
4081            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4082        }
4083        TapeOp::Erf(a) => {
4084            let u = vals[*a];
4085            let gp = erf_d1(u);
4086            let gpp = erf_d2(u);
4087            adj[*a] += w * gp;
4088            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4089        }
4090        TapeOp::XLogX(a) => {
4091            // See the identical arm in `hessian_accumulate`: gpp is 1/u, never
4092            // the unrepresentable -1/u² of `ln''`.
4093            let u = vals[*a];
4094            let gp = xlogx_d1(u);
4095            let gpp = xlogx_d2(u);
4096            adj[*a] += w * gp;
4097            adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4098        }
4099        TapeOp::CEntropy(a, b) => {
4100            let ua = vals[*a];
4101            let ub = vals[*b];
4102            let fa = centropy_da(ua, ub);
4103            let fb = centropy_db(ua, ub);
4104            let faa = centropy_daa(ua);
4105            let fab = centropy_dab(ub);
4106            let fbb = centropy_dbb(ua, ub);
4107            adj[*a] += w * fa;
4108            adj[*b] += w * fb;
4109            adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
4110            adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
4111        }
4112        TapeOp::Atan2(a, b) => {
4113            let y = vals[*a];
4114            let x = vals[*b];
4115            let (fa, fb) = atan2_d1(y, x);
4116            let (faa, fab, fbb) = atan2_d2(y, x);
4117            adj[*a] += w * fa;
4118            adj[*b] += w * fb;
4119            adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
4120            adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
4121        }
4122        TapeOp::Cmp(_, _, _)
4123        | TapeOp::And(_, _)
4124        | TapeOp::Or(_, _)
4125        | TapeOp::Not(_)
4126        | TapeOp::Select(_, _, _)
4127        | TapeOp::Min(_, _)
4128        | TapeOp::Max(_, _) => panic!(
4129            "GlobalTape free-function kernels do not implement conditional / logical \
4130             / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
4131             instead."
4132        ),
4133        TapeOp::Funcall(fc) => {
4134            let FuncallData { lib, name, args } = fc.as_ref();
4135            let call_args = funcall_to_ext_args(args, vals);
4136            let res = lib.eval(name, &call_args, true, true).unwrap_or_else(|e| {
4137                panic!("external function '{name}' 2nd-order eval failed: {e}")
4138            });
4139            let derivs = res.derivs.expect("want_derivs=true returns derivs");
4140            let hes = res.hessian.expect("want_hes=true returns hessian");
4141            let real_tape: Vec<usize> = args
4142                .iter()
4143                .filter_map(|a| match a {
4144                    TapeFuncallArg::Tape(t) => Some(*t),
4145                    TapeFuncallArg::Str(_) => None,
4146                })
4147                .collect();
4148            for (k, &tk) in real_tape.iter().enumerate() {
4149                adj[tk] += w * derivs[k];
4150                let mut second_term = 0.0;
4151                for (l, &tl) in real_tape.iter().enumerate() {
4152                    let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
4153                    let h_kl = hes[lo + hi * (hi + 1) / 2];
4154                    second_term += h_kl * dot[tl];
4155                }
4156                adj_dot[tk] += wd * derivs[k] + w * second_term;
4157            }
4158            let _ = out;
4159            let _ = weight;
4160            let _ = i;
4161        }
4162    }
4163}
4164
4165/// Per-op Hessian-sparsity propagation. Same algorithm as
4166/// `Tape::hessian_sparsity` but as a free function so `GlobalTape`
4167/// can call it over its shared `ops` slice.
4168fn hessian_sparsity_impl(ops: &[TapeOp]) -> BTreeSet<(usize, usize)> {
4169    let n = ops.len();
4170    let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
4171    let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
4172
4173    let emit_cross =
4174        |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4175            for &v1 in s1 {
4176                for &v2 in s2 {
4177                    let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
4178                    pairs.insert((r, c));
4179                }
4180            }
4181        };
4182    let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4183        let vars: Vec<usize> = s.iter().copied().collect();
4184        for (ai, &vi) in vars.iter().enumerate() {
4185            for &vj in &vars[..=ai] {
4186                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4187                pairs.insert((r, c));
4188            }
4189        }
4190    };
4191
4192    for op in ops {
4193        let vset = match op {
4194            TapeOp::Const(_) => BTreeSet::new(),
4195            TapeOp::Var(j) => {
4196                let mut s = BTreeSet::new();
4197                s.insert(*j);
4198                s
4199            }
4200            TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
4201                var_sets[*a].union(&var_sets[*b]).copied().collect()
4202            }
4203            TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
4204            TapeOp::Mul(a, b) => {
4205                emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4206                var_sets[*a].union(&var_sets[*b]).copied().collect()
4207            }
4208            TapeOp::Div(a, b) => {
4209                emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4210                emit_self(&var_sets[*b], &mut pairs);
4211                var_sets[*a].union(&var_sets[*b]).copied().collect()
4212            }
4213            TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
4214                let combined: BTreeSet<usize> =
4215                    var_sets[*a].union(&var_sets[*b]).copied().collect();
4216                emit_self(&combined, &mut pairs);
4217                combined
4218            }
4219            TapeOp::Sqrt(a)
4220            | TapeOp::Exp(a)
4221            | TapeOp::Log(a)
4222            | TapeOp::Log10(a)
4223            | TapeOp::Sin(a)
4224            | TapeOp::Cos(a)
4225            | TapeOp::Tan(a)
4226            | TapeOp::Atan(a)
4227            | TapeOp::Acos(a)
4228            | TapeOp::Sinh(a)
4229            | TapeOp::Cosh(a)
4230            | TapeOp::Tanh(a)
4231            | TapeOp::Asin(a)
4232            | TapeOp::Acosh(a)
4233            | TapeOp::Asinh(a)
4234            | TapeOp::Erf(a)
4235            | TapeOp::XLogX(a)
4236            | TapeOp::Atanh(a) => {
4237                emit_self(&var_sets[*a], &mut pairs);
4238                var_sets[*a].clone()
4239            }
4240            TapeOp::Funcall(fc) => {
4241                let args = &fc.args;
4242                let mut combined: BTreeSet<usize> = BTreeSet::new();
4243                for arg in args {
4244                    if let TapeFuncallArg::Tape(t) = arg {
4245                        for &vv in &var_sets[*t] {
4246                            combined.insert(vv);
4247                        }
4248                    }
4249                }
4250                emit_self(&combined, &mut pairs);
4251                combined
4252            }
4253            TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
4254                // Comparisons / logical ops have identically-zero derivative, so
4255                // they contribute no Hessian structure.
4256                BTreeSet::new()
4257            }
4258            TapeOp::Select(_, t, e) => {
4259                // Either branch may be active; the structural superset is the
4260                // union of both branches' variable sets.
4261                var_sets[*t].union(&var_sets[*e]).copied().collect()
4262            }
4263            TapeOp::Min(a, b) | TapeOp::Max(a, b) => {
4264                // min/max are piecewise linear: zero second derivative (no
4265                // pairs); dependence set is the union of both operands.
4266                var_sets[*a].union(&var_sets[*b]).copied().collect()
4267            }
4268        };
4269        var_sets.push(vset);
4270    }
4271    pairs
4272}
4273
4274#[cfg(test)]
4275mod tests {
4276    use super::*;
4277
4278    fn cnst(c: f64) -> Expr {
4279        Expr::Const(c)
4280    }
4281    fn var(i: usize) -> Expr {
4282        Expr::Var(i)
4283    }
4284    fn add(a: Expr, b: Expr) -> Expr {
4285        Expr::Binary(BinOp::Add, Box::new(a), Box::new(b))
4286    }
4287    fn mul(a: Expr, b: Expr) -> Expr {
4288        Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b))
4289    }
4290    fn pow(a: Expr, b: Expr) -> Expr {
4291        Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b))
4292    }
4293    fn div(a: Expr, b: Expr) -> Expr {
4294        Expr::Binary(BinOp::Div, Box::new(a), Box::new(b))
4295    }
4296    fn unary(op: UnaryOp, a: Expr) -> Expr {
4297        Expr::Unary(op, Box::new(a))
4298    }
4299    fn cmp(op: CmpOp, a: Expr, b: Expr) -> Expr {
4300        Expr::Compare(op, Box::new(a), Box::new(b))
4301    }
4302    fn cond(c: Expr, t: Expr, e: Expr) -> Expr {
4303        Expr::Cond {
4304            cond: Box::new(c),
4305            then_: Box::new(t),
4306            else_: Box::new(e),
4307        }
4308    }
4309
4310    #[test]
4311    fn polynomial_eval_and_grad() {
4312        // f = 3*x0^2 + 2*x1
4313        let e = add(
4314            mul(cnst(3.0), pow(var(0), cnst(2.0))),
4315            mul(cnst(2.0), var(1)),
4316        );
4317        let t = Tape::build(&e);
4318        assert!((t.eval(&[2.0, 3.0]) - 18.0).abs() < 1e-12);
4319        let mut g = vec![0.0; 2];
4320        t.gradient_seed(&[2.0, 3.0], 1.0, &mut g);
4321        // df/dx0 = 6*x0 = 12, df/dx1 = 2
4322        assert!((g[0] - 12.0).abs() < 1e-12);
4323        assert!((g[1] - 2.0).abs() < 1e-12);
4324    }
4325
4326    #[test]
4327    fn cse_shared_body_evaluated_once() {
4328        // body = x0 + x1, shared via Arc. f = body^2 + body.
4329        let body = Arc::new(add(var(0), var(1)));
4330        let e = add(
4331            pow(Expr::Cse(body.clone()), cnst(2.0)),
4332            Expr::Cse(body.clone()),
4333        );
4334        let t = Tape::build(&e);
4335        // body should appear once in the tape: count Add(Var(0),Var(1)) ops
4336        let n_body_adds = t
4337            .ops
4338            .iter()
4339            .filter(|op| {
4340                matches!(op, TapeOp::Add(a, b) if {
4341                    matches!(t.ops[*a], TapeOp::Var(0)) && matches!(t.ops[*b], TapeOp::Var(1))
4342                })
4343            })
4344            .count();
4345        assert_eq!(n_body_adds, 1, "CSE body should be emitted exactly once");
4346
4347        // f(1, 2) = 9 + 3 = 12
4348        assert!((t.eval(&[1.0, 2.0]) - 12.0).abs() < 1e-12);
4349        let mut g = vec![0.0; 2];
4350        t.gradient_seed(&[1.0, 2.0], 1.0, &mut g);
4351        // df/dx0 = 2*(x0+x1) + 1 = 7, same for x1
4352        assert!((g[0] - 7.0).abs() < 1e-12);
4353        assert!((g[1] - 7.0).abs() < 1e-12);
4354    }
4355
4356    fn fd_check(tape: &Tape, x: &[f64], n: usize, tol: f64) {
4357        let vars = tape.variables();
4358        let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4359        let mut pairs = Vec::new();
4360        for (ai, &vi) in vars.iter().enumerate() {
4361            for &vj in &vars[..=ai] {
4362                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4363                hess_map.entry((r, c)).or_insert_with(|| {
4364                    let p = pairs.len();
4365                    pairs.push((r, c));
4366                    p
4367                });
4368            }
4369        }
4370        let nnz = pairs.len();
4371        let mut ad = vec![0.0; nnz];
4372        tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
4373
4374        let mut fd = vec![0.0; nnz];
4375        let mut xp = x.to_vec();
4376        let mut gp = vec![0.0; n];
4377        let mut gm = vec![0.0; n];
4378        for &j in &vars {
4379            let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4380            xp[j] = x[j] + h;
4381            gp.iter_mut().for_each(|v| *v = 0.0);
4382            tape.gradient_seed(&xp, 1.0, &mut gp);
4383            xp[j] = x[j] - h;
4384            gm.iter_mut().for_each(|v| *v = 0.0);
4385            tape.gradient_seed(&xp, 1.0, &mut gm);
4386            xp[j] = x[j];
4387            for &i in &vars {
4388                if i >= j {
4389                    if let Some(&pos) = hess_map.get(&(i, j)) {
4390                        fd[pos] = (gp[i] - gm[i]) / (2.0 * h);
4391                    }
4392                }
4393            }
4394        }
4395        for (k, &(r, c)) in pairs.iter().enumerate() {
4396            let scale = fd[k].abs().max(1.0);
4397            assert!(
4398                (ad[k] - fd[k]).abs() / scale < tol,
4399                "H[{},{}]: AD={:.6e} FD={:.6e}",
4400                r,
4401                c,
4402                ad[k],
4403                fd[k]
4404            );
4405        }
4406    }
4407
4408    #[test]
4409    fn hessian_quadratic_matches_fd() {
4410        // f = 3 x0^2 + 2 x0 x1 + x1^2
4411        let e = add(
4412            add(
4413                mul(cnst(3.0), pow(var(0), cnst(2.0))),
4414                mul(cnst(2.0), mul(var(0), var(1))),
4415            ),
4416            pow(var(1), cnst(2.0)),
4417        );
4418        let t = Tape::build(&e);
4419        fd_check(&t, &[2.0, 3.0], 2, 1e-5);
4420    }
4421
4422    #[test]
4423    fn hessian_transcendental_matches_fd() {
4424        // f = exp(x0) + sin(x1) + log(x0) + sqrt(x1) + x0*x1
4425        let e = Expr::Sum(vec![
4426            unary(UnaryOp::Exp, var(0)),
4427            unary(UnaryOp::Sin, var(1)),
4428            unary(UnaryOp::Log, var(0)),
4429            unary(UnaryOp::Sqrt, var(1)),
4430            mul(var(0), var(1)),
4431        ]);
4432        let t = Tape::build(&e);
4433        fd_check(&t, &[1.5, 2.0], 2, 1e-5);
4434    }
4435
4436    #[test]
4437    fn inverse_trig_grad_and_hessian_match_fd() {
4438        // f = tan(x0) + atan(x1) + acos(x2) + x0*x1
4439        // Point chosen so every op is in its smooth domain:
4440        // tan away from pi/2, acos arg in (-1, 1).
4441        let e = Expr::Sum(vec![
4442            unary(UnaryOp::Tan, var(0)),
4443            unary(UnaryOp::Atan, var(1)),
4444            unary(UnaryOp::Acos, var(2)),
4445            mul(var(0), var(1)),
4446        ]);
4447        let t = Tape::build(&e);
4448        let x = [0.5, 1.3, 0.3];
4449
4450        // Gradient vs central finite difference of the value. This
4451        // pins the first derivatives independently of the Hessian
4452        // (fd_check only ties the Hessian to the AD gradient).
4453        let mut g = vec![0.0; 3];
4454        t.gradient_seed(&x, 1.0, &mut g);
4455        for j in 0..3 {
4456            let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4457            let mut xp = x;
4458            let mut xm = x;
4459            xp[j] += h;
4460            xm[j] -= h;
4461            let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4462            let scale = fd.abs().max(1.0);
4463            assert!(
4464                (g[j] - fd).abs() / scale < 1e-5,
4465                "grad[{j}]: AD={:.6e} FD={:.6e}",
4466                g[j],
4467                fd
4468            );
4469        }
4470
4471        // Hessian (forward-over-reverse) vs FD of the gradient.
4472        fd_check(&t, &x, 3, 1e-5);
4473    }
4474
4475    /// Shared helper: check AD gradient vs central FD of the value at
4476    /// `x`, then the Hessian via `fd_check`.
4477    fn grad_and_hess_match_fd(e: &Expr, x: &[f64], tol: f64) {
4478        let n = x.len();
4479        let t = Tape::build(e);
4480        let mut g = vec![0.0; n];
4481        t.gradient_seed(x, 1.0, &mut g);
4482        for j in 0..n {
4483            let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4484            let mut xp = x.to_vec();
4485            let mut xm = x.to_vec();
4486            xp[j] += h;
4487            xm[j] -= h;
4488            let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4489            let scale = fd.abs().max(1.0);
4490            assert!(
4491                (g[j] - fd).abs() / scale < tol,
4492                "grad[{j}]: AD={:.6e} FD={:.6e}",
4493                g[j],
4494                fd
4495            );
4496        }
4497        fd_check(&t, x, n, tol);
4498    }
4499
4500    #[test]
4501    fn hyperbolic_grad_and_hessian_match_fd() {
4502        // f = sinh(x0) + cosh(x1) + tanh(x2) + asinh(x3) + x0*x1 + x2*x3
4503        // sinh/cosh/tanh/asinh are smooth on all of R.
4504        let e = Expr::Sum(vec![
4505            unary(UnaryOp::Sinh, var(0)),
4506            unary(UnaryOp::Cosh, var(1)),
4507            unary(UnaryOp::Tanh, var(2)),
4508            unary(UnaryOp::Asinh, var(3)),
4509            mul(var(0), var(1)),
4510            mul(var(2), var(3)),
4511        ]);
4512        grad_and_hess_match_fd(&e, &[0.5, 0.7, 0.3, 1.1], 1e-5);
4513    }
4514
4515    #[test]
4516    fn restricted_inverse_grad_and_hessian_match_fd() {
4517        // f = asin(x0) + acosh(x1) + atanh(x2) + x0*x2
4518        // Point chosen in each op's smooth domain:
4519        // asin/atanh need |arg| < 1; acosh needs arg > 1.
4520        let e = Expr::Sum(vec![
4521            unary(UnaryOp::Asin, var(0)),
4522            unary(UnaryOp::Acosh, var(1)),
4523            unary(UnaryOp::Atanh, var(2)),
4524            mul(var(0), var(2)),
4525        ]);
4526        grad_and_hess_match_fd(&e, &[0.4, 1.8, 0.3], 1e-5);
4527    }
4528
4529    #[test]
4530    fn erf_value_matches_reference() {
4531        // Reference values from the standard erf (musl / C99), to full
4532        // double precision. Pinning these guards the `libm` delegation:
4533        // a swap to a series approximation would fail here long before it
4534        // showed up as a mysteriously loose KKT residual.
4535        let t = Tape::build(&unary(UnaryOp::Erf, var(0)));
4536        for (x, want) in [
4537            (0.0, 0.0),
4538            (0.5, 0.520_499_877_813_046_5),
4539            (1.0, 0.842_700_792_949_714_9),
4540            (-1.0, -0.842_700_792_949_714_9),
4541            (2.0, 0.995_322_265_018_952_7),
4542            (3.0, 0.999_977_909_503_001_4),
4543        ] {
4544            let got = t.eval(&[x]);
4545            assert!(
4546                (got - want).abs() < 1e-15,
4547                "erf({x}): got {got:.17e}, want {want:.17e}"
4548            );
4549        }
4550        // Odd and saturating: the two properties a wrong implementation
4551        // most often breaks.
4552        assert!((t.eval(&[-0.3]) + t.eval(&[0.3])).abs() < 1e-16);
4553        assert!((t.eval(&[10.0]) - 1.0).abs() < 1e-15);
4554    }
4555
4556    #[test]
4557    fn atanh_value_and_derivative_accurate_near_boundary() {
4558        // gh#1151: `f64::atanh` computes `0.5 * ln_1p(2u / (1 - u))`, which
4559        // loses `δ` at `u = -1 + δ` (rel. error 3.6e-5 at δ = 1e-13, 3.0e-3
4560        // at δ = 1e-15). The derivative's `1 - u*u` is accurate this close to
4561        // the boundary; its worst case is further in — see
4562        // `INVERSE_DERIV_CASES`.
4563        // References: `0.5·ln((1+u)/(1-u))` and `1/((1-u)(1+u))` in 60-digit
4564        // decimal at the exact f64 input. Both signs of the boundary are
4565        // pinned — the defect was asymmetric, so a symmetric fixture alone
4566        // would not reach the branch that broke.
4567        let t = Tape::build(&unary(UnaryOp::Atanh, var(0)));
4568        let rel = |got: f64, want: f64| ((got - want) / want).abs();
4569        for (u, want_v, want_d) in [
4570            (-0.9999999999999, -15.31322124631433, 4998445757348.192),
4571            (-0.999999999999999, -17.616361586450413, 500399958596722.0),
4572            (0.999999999999999, 17.616361586450413, 500399958596722.0),
4573            (0.9999999999999998, 18.36840028483855, 2251799813685248.5),
4574            (-0.9, -1.4722194895832204, 5.263157894736843),
4575            (0.5, 0.5493061443340549, 1.3333333333333333),
4576            (1e-10, 1e-10, 1.0),
4577        ] {
4578            let got_v = t.eval(&[u]);
4579            assert!(
4580                rel(got_v, want_v) < 4e-16,
4581                "atanh({u:e}): got {got_v:.17e}, want {want_v:.17e}"
4582            );
4583            let mut g = [0.0];
4584            t.gradient_seed(&[u], 1.0, &mut g);
4585            assert!(
4586                rel(g[0], want_d) < 4e-16,
4587                "atanh'({u:e}): got {:.17e}, want {want_d:.17e}",
4588                g[0]
4589            );
4590        }
4591    }
4592
4593    /// `got` matches `want` to `tol` relative, or both are zero. A subnormal
4594    /// `want` carries fewer than 53 bits, so there one subnormal ulp of
4595    /// absolute error is allowed as well.
4596    fn close(got: f64, want: f64, tol: f64) -> bool {
4597        if want == 0.0 {
4598            got == 0.0
4599        } else if want.abs() < f64::MIN_POSITIVE {
4600            (got - want).abs() <= f64::from_bits(1)
4601        } else {
4602            ((got - want) / want).abs() < tol
4603        }
4604    }
4605
4606    /// (op, u, f'(u), f''(u)), references from 100-digit decimal at the exact
4607    /// f64 input. Each op has an interior point (the textbook form was fine)
4608    /// and the points where the textbook form failed: cancellation near the
4609    /// boundary of the domain, or overflow of `u*u` / `d*d` at large `|u|`.
4610    ///
4611    /// The cancellation points sit at `δ ≈ 3e-9` from the boundary, not at
4612    /// `δ = 1e-15`: for `u = 1 - δ` the rounding error of `u*u` is ~`δ²`
4613    /// (until that reaches an ulp), so `1 - u*u` is accurate *at* the
4614    /// boundary and worst (~4e-9 relative) where `δ²` is one ulp. The `1e-15`
4615    /// points stay because they pin the fixed forms at the boundary itself.
4616    const INVERSE_DERIV_CASES: &[(UnaryOp, f64, f64, f64)] = &[
4617        (UnaryOp::Asin, 0.3, 1.0482848367219182, 0.3455884077105225),
4618        (
4619            UnaryOp::Asin,
4620            -0.999999999999999,
4621            22369621.33333334,
4622            -1.1193757589024341e22,
4623        ),
4624        (
4625            UnaryOp::Asin,
4626            0.999999999999999,
4627            22369621.33333334,
4628            1.1193757589024341e22,
4629        ),
4630        (
4631            UnaryOp::Asin,
4632            0.999999992554274,
4633            8194.670143394766,
4634            550293555898.1869,
4635        ),
4636        (UnaryOp::Acos, 0.3, -1.0482848367219182, -0.3455884077105225),
4637        (
4638            UnaryOp::Acos,
4639            -0.999999999999999,
4640            -22369621.33333334,
4641            1.1193757589024341e22,
4642        ),
4643        (
4644            UnaryOp::Acos,
4645            0.999999992554274,
4646            -8194.670143394766,
4647            -550293555898.1869,
4648        ),
4649        (UnaryOp::Acosh, 2.0, 0.5773502691896257, -0.3849001794597505),
4650        (
4651            UnaryOp::Acosh,
4652            1.0000000000001,
4653            2236962.1333332774,
4654            -1.1193757589024625e19,
4655        ),
4656        (
4657            UnaryOp::Acosh,
4658            1.0000000050104683,
4659            9989.54813402654,
4660            -996867721306.0283,
4661        ),
4662        (UnaryOp::Acosh, 1e100, 1e-100, -1e-200),
4663        (UnaryOp::Acosh, 1e200, 1e-200, 0.0),
4664        (
4665            UnaryOp::Asinh,
4666            0.5,
4667            0.8944271909999159,
4668            -0.35777087639996635,
4669        ),
4670        (UnaryOp::Asinh, 1e100, 1e-100, -1e-200),
4671        (UnaryOp::Asinh, -1e100, 1e-100, 1e-200),
4672        (UnaryOp::Asinh, 1e200, 1e-200, 0.0),
4673        (UnaryOp::Asinh, -1e200, 1e-200, 0.0),
4674        (UnaryOp::Atanh, 0.5, 1.3333333333333333, 1.7777777777777777),
4675        (
4676            UnaryOp::Atanh,
4677            -0.999999999999999,
4678            500399958596722.0,
4679            -5.008002371272017e29,
4680        ),
4681        (
4682            UnaryOp::Atanh,
4683            0.999999999999999,
4684            500399958596722.0,
4685            5.008002371272017e29,
4686        ),
4687        (
4688            UnaryOp::Atanh,
4689            -0.999999992554274,
4690            67152618.75904559,
4691            -9018948345242822.0,
4692        ),
4693        (UnaryOp::Atan, 0.7, 0.6711409395973155, -0.6306022251249944),
4694        (UnaryOp::Atan, 1e100, 1e-200, -2e-300),
4695        (UnaryOp::Atan, -1e100, 1e-200, 2e-300),
4696        // Past `|u| ≈ 1.3e154`, `1 + u*u` is `inf` and the old `atan'` read
4697        // `0`; the true value stays subnormal up to `|u| ≈ 6e161`. 2e150 is
4698        // just past the `r = 1/u` switch, where the value is still normal.
4699        (UnaryOp::Atan, 2e150, 2.5e-301, 0.0),
4700        (UnaryOp::Atan, 1e155, 1e-310, 0.0),
4701        (UnaryOp::Atan, -1e158, 1e-316, 0.0),
4702        (UnaryOp::Atan, 1e160, 1e-320, 0.0),
4703        (UnaryOp::Atan, 5e161, 5e-324, 0.0),
4704        (UnaryOp::Tanh, 0.3, 0.9151369618266292, -0.5331818782014544),
4705        (
4706            UnaryOp::Tanh,
4707            10.0,
4708            8.244614455767397e-9,
4709            -1.6489228843561127e-8,
4710        ),
4711        (
4712            UnaryOp::Tanh,
4713            25.0,
4714            7.714999391855671e-22,
4715            -1.5429998783711343e-21,
4716        ),
4717        (
4718            UnaryOp::Tanh,
4719            -30.0,
4720            3.502604305078608e-26,
4721            7.005208610157216e-26,
4722        ),
4723    ];
4724
4725    #[test]
4726    fn inverse_and_tanh_derivatives_accurate_at_boundaries_and_extremes() {
4727        // The neighbours of gh#1151: every one of these failed in the
4728        // textbook form — `1 - u*u` / `u*u - 1` cancelled at the domain
4729        // boundary, `1 - tanh²` cancelled once tanh rounded toward ±1 (and
4730        // was exactly 0 past |u| ≈ 19), and `u*u + 1` / `d*d` overflowed to
4731        // a zero derivative at large |u|. Checked through the tape — the
4732        // reverse gradient and `hessian_accumulate` — so the arms are
4733        // pinned, not just the kernels.
4734        assert_unary_derivs(INVERSE_DERIV_CASES);
4735    }
4736
4737    /// Checks each (op, u, f', f'') through the tape: the reverse gradient
4738    /// and `hessian_accumulate`, so the arms are pinned, not just the kernels.
4739    fn assert_unary_derivs(cases: &[(UnaryOp, f64, f64, f64)]) {
4740        for &(op, u, want_d1, want_d2) in cases {
4741            let t = Tape::build(&unary(op, var(0)));
4742            let mut g = [0.0];
4743            t.gradient_seed(&[u], 1.0, &mut g);
4744            assert!(
4745                close(g[0], want_d1, 1e-15),
4746                "{op:?}'({u:e}): got {:.17e}, want {want_d1:.17e}",
4747                g[0]
4748            );
4749            // Second derivatives get 2e-15: the kernels are a few roundings
4750            // deep (`acosh″` reaches 1.06e-15 on random inputs), and
4751            // `tanh`/`asinh`/`atan2` go through the platform's `cosh`/`hypot`,
4752            // which CI only runs on Linux.
4753            let mut hess_map = HashMap::new();
4754            hess_map.insert((0usize, 0usize), 0usize);
4755            let mut h = [0.0];
4756            t.hessian_accumulate(&[u], 1.0, &hess_map, &mut h);
4757            assert!(
4758                close(h[0], want_d2, 2e-15),
4759                "{op:?}''({u:e}): got {:.17e}, want {want_d2:.17e}",
4760                h[0]
4761            );
4762        }
4763    }
4764
4765    /// The same overflow class in `log`/`log10`/`sqrt`: `-1/(u*u)`,
4766    /// `1/(u·ln 10)` and `-0.25/(u·sqrt u)` read `0` once the product
4767    /// overflowed, while the true value is still representable (normal or
4768    /// subnormal). Each op has points on both sides of its switch, since the
4769    /// old form is kept bit for bit below it and the rescaled form takes over
4770    /// above: `log''`/`log10''` at 1e150, `log10'` at 1e300, `sqrt''` at 1e200.
4771    const LOG_SQRT_DERIV_CASES: &[(UnaryOp, f64, f64, f64)] = &[
4772        (UnaryOp::Log, 2.0, 0.5, -0.25),
4773        (UnaryOp::Log, 1e100, 1e-100, -1e-200),
4774        (UnaryOp::Log, 1e155, 1e-155, -1e-310),
4775        (UnaryOp::Log, 1e160, 1e-160, -1e-320),
4776        (
4777            UnaryOp::Log10,
4778            2.0,
4779            0.2171472409516259,
4780            -0.10857362047581295,
4781        ),
4782        (
4783            UnaryOp::Log10,
4784            1e100,
4785            4.342944819032518e-101,
4786            -4.342944819032518e-201,
4787        ),
4788        (
4789            UnaryOp::Log10,
4790            1e155,
4791            4.3429448190325184e-156,
4792            -4.3429448190326e-311,
4793        ),
4794        (UnaryOp::Log10, 1e305, 4.3429448190325186e-306, 0.0),
4795        (UnaryOp::Log10, 1.7e308, 2.554673422960305e-309, 0.0),
4796        (UnaryOp::Sqrt, 2.0, 0.3535533905932738, -0.08838834764831845),
4797        (UnaryOp::Sqrt, 1e150, 5e-76, -2.5e-226),
4798        (UnaryOp::Sqrt, 1e206, 5e-104, -2.5e-310),
4799        (UnaryOp::Sqrt, 1e300, 5e-151, 0.0),
4800        (UnaryOp::Sqrt, 1.7e308, 3.8348249442368524e-155, 0.0),
4801    ];
4802
4803    #[test]
4804    fn log_log10_sqrt_derivatives_accurate_at_large_arguments() {
4805        assert_unary_derivs(LOG_SQRT_DERIV_CASES);
4806    }
4807
4808    #[test]
4809    fn atan2_derivatives_accurate_at_extreme_magnitudes() {
4810        // `d = y*y + x*x` overflowed past ~1e154 (zero gradient) and
4811        // underflowed below ~1e-162 (infinite gradient); `d*d` in the
4812        // Hessian did both at the square root of those magnitudes. The last
4813        // case has |y| ≈ |x|, where an unfactored `y² - x²` cancels.
4814        let e = Expr::Binary(BinOp::Atan2, Box::new(var(0)), Box::new(var(1)));
4815        let t = Tape::build(&e);
4816        let mut hess_map = HashMap::new();
4817        hess_map.insert((0usize, 0usize), 0usize);
4818        hess_map.insert((1usize, 0usize), 1usize);
4819        hess_map.insert((1usize, 1usize), 2usize);
4820        for (y, x, want_g, want_h) in [
4821            (
4822                1e100,
4823                2e100,
4824                [4e-101, -2e-101],
4825                [-1.6e-201, -1.2e-201, 1.6e-201],
4826            ),
4827            (1e-100, -3e-100, [-3e99, -1e99], [6e198, -8e198, -6e198]),
4828            (1.0, 2.0, [0.4, -0.2], [-0.16, -0.12, 0.16]),
4829            // `y ± x` overflows here: the old off-diagonal was `NaN` at
4830            // `y = -x` (`inf · 0`) and `inf` beside it. The true Hessian
4831            // underflows to `0`; the gradient is subnormal.
4832            (1e308, -1e308, [-5e-309, -5e-309], [0.0, 0.0, 0.0]),
4833            (
4834                1.5e308,
4835                -1e308,
4836                [-3.076923076923076e-309, -4.615384615384614e-309],
4837                [0.0, 0.0, 0.0],
4838            ),
4839            (1e308, 1e308, [5e-309, -5e-309], [0.0, 0.0, 0.0]),
4840            (
4841                3.0,
4842                3.0000000001,
4843                [0.16666666666666666, -0.1666666666611111],
4844                [
4845                    -0.0555555555537037,
4846                    -1.8518520049821686e-12,
4847                    0.0555555555537037,
4848                ],
4849            ),
4850        ] {
4851            let mut g = [0.0; 2];
4852            t.gradient_seed(&[y, x], 1.0, &mut g);
4853            let mut h = [0.0; 3];
4854            t.hessian_accumulate(&[y, x], 1.0, &hess_map, &mut h);
4855            for k in 0..2 {
4856                assert!(
4857                    close(g[k], want_g[k], 1e-15),
4858                    "atan2({y:e},{x:e}) grad[{k}]: got {:.17e}, want {:.17e}",
4859                    g[k],
4860                    want_g[k]
4861                );
4862            }
4863            for k in 0..3 {
4864                assert!(
4865                    close(h[k], want_h[k], 2e-15),
4866                    "atan2({y:e},{x:e}) hess[{k}]: got {:.17e}, want {:.17e}",
4867                    h[k],
4868                    want_h[k]
4869                );
4870            }
4871        }
4872    }
4873
4874    #[test]
4875    fn erf_second_derivative_stays_finite_at_extreme_magnitudes() {
4876        // `erf_d1` underflows to 0 around |u| > 27, and `-2.0 * u`
4877        // overflows to ±inf past f64::MAX/2. Written as `(-2u)·erf'(u)`
4878        // the two meet as `inf * 0.0 = NaN`, where the true limit is 0.
4879        // Parenthesizing as `-2·(u·erf'(u))` keeps every finite input
4880        // finite. Nothing in a well-posed model reaches 1e308 — `.nl`'s
4881        // unbounded sentinel is 1e19 — but the Python binding now takes an
4882        // unguarded `x`, and a NaN Hessian entry poisons a whole factorization.
4883        for u in [1e19, 1e150, 1e300, f64::MAX, f64::MAX / 2.0] {
4884            for signed in [u, -u] {
4885                let d2 = erf_d2(signed);
4886                assert!(
4887                    d2.is_finite(),
4888                    "erf_d2({signed:e}) = {d2} — must be finite (the limit is 0)"
4889                );
4890            }
4891        }
4892        // Still correct where it matters: -2u·erf'(u) at a normal point.
4893        let u = 0.7;
4894        let want = -2.0 * u * (2.0 / std::f64::consts::PI.sqrt()) * (-u * u).exp();
4895        assert!((erf_d2(u) - want).abs() < 1e-15, "{}", erf_d2(u));
4896    }
4897
4898    #[test]
4899    fn quotient_rule_keeps_representable_derivatives_at_extreme_denominators() {
4900        // The quotient rule used to form `b*b` explicitly. That overflows for
4901        // |b| > 1.3e154 and underflows below 1.5e-154, and what it destroys is a
4902        // derivative that is itself perfectly representable: at a = b = 1e300,
4903        // ∂(a/b)/∂b = -a/b² = -1e-300, but `b*b` is inf so the squared form
4904        // returned -0.0. A silently ZEROED gradient entry, not a NaN anyone
4905        // would catch downstream.
4906        //
4907        // Reachable in practice: discopt lowers GAMS `centropy(x, y)` to
4908        // `x·log(x/y)`, and a 600-digit audit measured a wrong y-derivative at
4909        // (1e300, 1e300) through exactly this path.
4910        let t = Tape::build(&div(var(0), var(1)));
4911        let mut g = vec![0.0; 2];
4912        for b in [1e300_f64, 1e200, 1e160, -1e300, 1e-160, 1e-300] {
4913            // At a = b: ∂/∂a = 1/b and ∂/∂b = -a/b² = -1/b, both representable
4914            // for every b here even though b² is not.
4915            // `gradient_seed` ACCUMULATES; without this reset the previous
4916            // iteration's entry survives and the assert compares against stale
4917            // data (it did, and caught itself: 1e-160 leaked into the b=-1e300 row).
4918            g.iter_mut().for_each(|v| *v = 0.0);
4919            t.gradient_seed(&[b, b], 1.0, &mut g);
4920            assert_eq!(g[0], 1.0 / b, "d(a/b)/da at a=b={b:e}");
4921            assert_eq!(
4922                g[1],
4923                -1.0 / b,
4924                "d(a/b)/db at a=b={b:e} — b² is not representable"
4925            );
4926        }
4927
4928        // Second order, where the rule squared and CUBED the denominator.
4929        // b = 1e155 makes b² = 1e310 overflow while the cross partial
4930        // ∂²(a/b)/∂a∂b = -1/b² = -1e-310 is still (subnormally) representable.
4931        let b = 1e155_f64;
4932        let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4933        hess_map.insert((0, 0), 0);
4934        hess_map.insert((1, 0), 1);
4935        hess_map.insert((1, 1), 2);
4936        let mut ad = vec![0.0; 3];
4937        t.hessian_accumulate(&[b, b], 1.0, &hess_map, &mut ad);
4938        // Compute the truth as -(1/b)/b, NOT -1/(b*b) -- the latter is the very
4939        // expression under test and evaluates to -0.0 here, which would have made
4940        // this assertion compare the bug against itself and pass vacuously.
4941        let want = -(1.0 / b) / b;
4942        assert!(
4943            want != 0.0 && want.is_finite(),
4944            "test premise: -1/b² must be representable, got {want:e}"
4945        );
4946        assert!(
4947            (ad[1] - want).abs() <= 1e-10 * want.abs(),
4948            "d²(a/b)/da db at a=b={b:e}: got {:e}, want {want:e}",
4949            ad[1]
4950        );
4951    }
4952
4953    #[test]
4954    fn erf_grad_and_hessian_match_fd() {
4955        // f = erf(x0) + erf(2*x1) + x0*x1. The inner `2*x1` makes the
4956        // chain rule non-trivial, so a missing factor in the tangent or
4957        // adjoint arm shows up rather than cancelling.
4958        let e = Expr::Sum(vec![
4959            unary(UnaryOp::Erf, var(0)),
4960            unary(UnaryOp::Erf, mul(cnst(2.0), var(1))),
4961            mul(var(0), var(1)),
4962        ]);
4963        grad_and_hess_match_fd(&e, &[0.4, -0.7], 1e-5);
4964    }
4965
4966    #[test]
4967    fn erf_directional_hessian_matches_accumulated() {
4968        // `hessian_directional` (the HVP / coloring sweep) and
4969        // `hessian_accumulate` (the sparse-entry sweep) are separate
4970        // match arms over the same op — the easiest place for an erf
4971        // second derivative to be right in one and wrong in the other.
4972        let e = Expr::Sum(vec![
4973            unary(UnaryOp::Erf, var(0)),
4974            unary(UnaryOp::Erf, mul(var(0), var(1))),
4975        ]);
4976        let tape = Tape::build(&e);
4977        let x = [0.6, -0.9];
4978        let n = x.len();
4979
4980        let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
4981        let hess_map: HashMap<(usize, usize), usize> =
4982            pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
4983        let mut acc = vec![0.0; pairs.len()];
4984        tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
4985
4986        // One directional product per unit seed reproduces column j.
4987        let ops = tape.ops.len();
4988        let mut vals = vec![0.0; ops];
4989        tape.forward_into(&x, &mut vals);
4990        for j in 0..n {
4991            let mut seed = vec![0.0; n];
4992            seed[j] = 1.0;
4993            let mut col = vec![0.0; n];
4994            let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
4995            tape.hessian_directional(
4996                &vals,
4997                &seed,
4998                1.0,
4999                &mut col,
5000                &mut dot,
5001                &mut adj,
5002                &mut adj_dot,
5003            );
5004            for i in 0..n {
5005                let (r, c) = if i >= j { (i, j) } else { (j, i) };
5006                let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
5007                assert!(
5008                    (col[i] - want).abs() < 1e-12,
5009                    "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
5010                    col[i]
5011                );
5012            }
5013        }
5014    }
5015
5016    fn centropy_expr(a: Expr, b: Expr) -> Expr {
5017        Expr::Binary(BinOp::CEntropy, Box::new(a), Box::new(b))
5018    }
5019
5020    /// Dense 2×2 Hessian via `hessian_accumulate`, as `(h00, h10, h11)`.
5021    fn hess2(e: &Expr, x: &[f64; 2]) -> (f64, f64, f64) {
5022        let tape = Tape::build(e);
5023        let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5024        hess_map.insert((0, 0), 0);
5025        hess_map.insert((1, 0), 1);
5026        hess_map.insert((1, 1), 2);
5027        let mut ad = vec![0.0; 3];
5028        tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
5029        (ad[0], ad[1], ad[2])
5030    }
5031
5032    #[test]
5033    fn xlogx_and_centropy_grad_and_hessian_match_fd() {
5034        // f = xlogx(x0) + centropy(x0, x1) + x0*x1. The trailing product
5035        // makes the cross partial nonzero, so a dropped ∂²/∂a∂b term shows
5036        // up instead of hiding in an all-zero off-diagonal.
5037        let e = Expr::Sum(vec![
5038            unary(UnaryOp::XLogX, var(0)),
5039            centropy_expr(var(0), var(1)),
5040            mul(var(0), var(1)),
5041        ]);
5042        grad_and_hess_match_fd(&e, &[1.7, 0.6], 1e-5);
5043    }
5044
5045    #[test]
5046    fn centropy_directional_hessian_matches_accumulated() {
5047        // `hessian_directional` and `hessian_accumulate` are separate match
5048        // arms over the same op — the easiest place for a second derivative
5049        // to be right in one sweep and wrong in the other.
5050        let e = Expr::Sum(vec![
5051            centropy_expr(var(0), var(1)),
5052            unary(UnaryOp::XLogX, mul(var(0), var(1))),
5053        ]);
5054        let tape = Tape::build(&e);
5055        let x = [1.3, 0.8];
5056        let n = x.len();
5057
5058        let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
5059        let hess_map: HashMap<(usize, usize), usize> =
5060            pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
5061        let mut acc = vec![0.0; pairs.len()];
5062        tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
5063
5064        let ops = tape.ops.len();
5065        let mut vals = vec![0.0; ops];
5066        tape.forward_into(&x, &mut vals);
5067        let mut compared = 0usize;
5068        for j in 0..n {
5069            let mut seed = vec![0.0; n];
5070            seed[j] = 1.0;
5071            let mut col = vec![0.0; n];
5072            let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
5073            tape.hessian_directional(
5074                &vals,
5075                &seed,
5076                1.0,
5077                &mut col,
5078                &mut dot,
5079                &mut adj,
5080                &mut adj_dot,
5081            );
5082            for i in 0..n {
5083                let (r, c) = if i >= j { (i, j) } else { (j, i) };
5084                let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
5085                assert!(
5086                    (col[i] - want).abs() < 1e-12,
5087                    "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
5088                    col[i]
5089                );
5090                compared += 1;
5091            }
5092        }
5093        // §6: prove the probe fired. A tape whose sparsity came back empty
5094        // would sail through the loop above comparing 0.0 against 0.0.
5095        assert_eq!(compared, 4, "expected a full 2x2 comparison");
5096        assert!(
5097            acc.iter().any(|v| v.abs() > 1e-12),
5098            "every accumulated Hessian entry was zero — the arms under test never ran"
5099        );
5100    }
5101
5102    #[test]
5103    fn fused_entropy_ops_reach_derivatives_the_chain_rule_cannot() {
5104        // The reason these ops exist. `(a·ln a)'' = 1/a` is finite for every
5105        // positive a down to 1e-308, but EVERY decomposition of it routes
5106        // through `ln''(a) = -1/a²`, which leaves f64 range below a ≈ 1e-154.
5107        // The composite is in range while the factor it is built from is not,
5108        // so this is a structural limit of the chain rule, not a sloppy rule.
5109        let a = 1e-299_f64;
5110        let want = 1.0 / a; // 1e299 — an ordinary number.
5111        assert!(want.is_finite(), "test premise: 1/a must be representable");
5112
5113        let (fused, _, _) = hess2(&unary(UnaryOp::XLogX, var(0)), &[a, 0.0]);
5114        assert!(
5115            (fused - want).abs() <= 1e-12 * want,
5116            "xlogx''({a:e}): got {fused:e}, want {want:e}"
5117        );
5118
5119        // The control: the same function written as `a * log(a)`. Asserting it
5120        // is BROKEN is what proves the test is in the regime the fusion exists
5121        // for — without this, a passing `fused` assertion could just mean the
5122        // point was never extreme enough to matter.
5123        let (decomposed, _, _) = hess2(&mul(var(0), unary(UnaryOp::Log, var(0))), &[a, 0.0]);
5124        assert!(
5125            !decomposed.is_finite(),
5126            "test premise: a*log(a) is supposed to lose this second derivative, \
5127             but it returned {decomposed:e} — if the Log/Mul arms now reach it, \
5128             re-derive whether the fused op is still needed"
5129        );
5130
5131        // Same story for centropy's ∂²/∂a², which is also 1/a.
5132        let (fused_aa, _, _) = hess2(&centropy_expr(var(0), var(1)), &[a, 1.0]);
5133        assert!(
5134            (fused_aa - want).abs() <= 1e-12 * want,
5135            "centropy ∂²/∂a² at a={a:e}: got {fused_aa:e}, want {want:e}"
5136        );
5137    }
5138
5139    #[test]
5140    fn centropy_second_derivative_in_b_survives_a_squared_denominator() {
5141        // ∂²/∂b² = a/b². At (1e300, 1e200) the answer is 1e-100 while b² is
5142        // 1e400 — out of range. Evaluating it as (a/b)/b keeps it.
5143        let (a, b) = (1e300_f64, 1e200_f64);
5144        let want = (a / b) / b;
5145        // The squared form does not blow up loudly — `b*b` is inf and
5146        // `a/inf` is a silent 0.0, i.e. a ZEROED second derivative. Assert
5147        // that shape explicitly so the premise says what actually goes wrong.
5148        assert!(
5149            want.is_finite() && want != 0.0 && a / (b * b) == 0.0,
5150            "test premise: a/b² representable ({want:e}) but a/(b*b) collapses"
5151        );
5152        let (_, _, h11) = hess2(&centropy_expr(var(0), var(1)), &[a, b]);
5153        assert!(
5154            (h11 - want).abs() <= 1e-12 * want,
5155            "centropy ∂²/∂b² at ({a:e}, {b:e}): got {h11:e}, want {want:e}"
5156        );
5157    }
5158
5159    #[test]
5160    fn centropy_value_survives_an_out_of_range_ratio_and_the_zero_limit() {
5161        let t = |e: &Expr, x: &[f64]| Tape::build(e).eval(x);
5162        let ce = centropy_expr(var(0), var(1));
5163
5164        // a/b = 1e600 overflows, but ln of it is an ordinary 1381.55, so the
5165        // product is a perfectly representable 1.38e303. `ln_ratio` falls back
5166        // to ln(a) - ln(b) exactly here.
5167        let want = 1e300 * (1e300_f64.ln() - 1e-300_f64.ln());
5168        let got = t(&ce, &[1e300, 1e-300]);
5169        assert!(
5170            got.is_finite() && (got - want).abs() <= 1e-12 * want,
5171            "centropy(1e300, 1e-300): got {got:e}, want {want:e}"
5172        );
5173        // The naive form is what this replaces.
5174        assert!(
5175            !t(
5176                &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
5177                &[1e300, 1e-300]
5178            )
5179            .is_finite(),
5180            "test premise: a*log(a/b) is supposed to overflow here"
5181        );
5182
5183        // Near a = b, where forming the ratio first throws away the digits
5184        // that carry the answer. At a = 1e10+1, b = 1e10 the quotient rounds
5185        // to 1 + 1e-10 with ~1e-16 absolute slop — 1e-6 RELATIVE slop in the
5186        // part that survives ln — while (a - b) is exactly 1.0 by Sterbenz.
5187        // Taylor: a·ln(1+1e-10) = 1 + 5e-11 - 1.7e-21, so 1.00000000005 is
5188        // the correctly-rounded f64 (the neglected term is ~1e-5 ulp).
5189        let (a, b) = (1e10 + 1.0, 1e10_f64);
5190        assert_eq!(a - b, 1.0, "test premise: a - b must be exact here");
5191        let got = t(&ce, &[a, b]);
5192        assert!(
5193            (got - 1.00000000005).abs() <= 2.3e-16,
5194            "centropy(1e10+1, 1e10): got {got:.17e}, want 1.00000000005"
5195        );
5196        // The naive form is off by 8e-8 here — eight digits, not one ulp.
5197        let naive = t(
5198            &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
5199            &[a, b],
5200        );
5201        assert!(
5202            (naive - 1.00000000005).abs() > 1e-9,
5203            "test premise: a*log(a/b) is supposed to lose digits here, got {naive:.17e}"
5204        );
5205
5206        // 0·ln(0/b) = 0 by the limit, not NaN.
5207        assert_eq!(t(&ce, &[0.0, 2.0]), 0.0);
5208        assert_eq!(t(&unary(UnaryOp::XLogX, var(0)), &[0.0]), 0.0);
5209    }
5210
5211    #[test]
5212    fn cond_does_not_leak_a_non_finite_from_its_inactive_branch() {
5213        // Frontends clamp `entropy`/`centropy` with a Cond so the x -> 0+ limit
5214        // stays finite: `select(x < 1e-300, x*ln(1e-300), xlogx(x))`. The tape's
5215        // forward sweep evaluates EVERY slot, so at x = -1 the inactive
5216        // `xlogx(-1)` arm is NaN and its derivative arms are NaN too. Whether
5217        // that NaN reaches the caller is the difference between a usable clamp
5218        // and a poisoned Hessian row, and it is decided independently in three
5219        // sweeps — so check all three, not just the gradient.
5220        let cond = Expr::Compare(CmpOp::Lt, Box::new(var(0)), Box::new(cnst(1e-300)));
5221        let e = Expr::Cond {
5222            cond: Box::new(cond),
5223            then_: Box::new(mul(var(0), cnst(1e-300f64.ln()))),
5224            else_: Box::new(unary(UnaryOp::XLogX, var(0))),
5225        };
5226        let t = Tape::build(&e);
5227        let x = [-1.0_f64];
5228
5229        // Premise: the inactive arm really is NaN at this point, so the test is
5230        // exercising the leak it claims to rule out.
5231        assert!(
5232            Tape::build(&unary(UnaryOp::XLogX, var(0)))
5233                .eval(&x)
5234                .is_nan(),
5235            "test premise: xlogx(-1) must be NaN"
5236        );
5237
5238        assert_eq!(t.eval(&x), -1e-300_f64.ln());
5239
5240        let mut g = vec![0.0; 1];
5241        t.gradient_seed(&x, 1.0, &mut g);
5242        assert_eq!(g[0], 1e-300_f64.ln(), "gradient leaked the inactive branch");
5243
5244        let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5245        hess_map.insert((0, 0), 0);
5246        let mut acc = vec![0.0; 1];
5247        t.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
5248        assert_eq!(acc[0], 0.0, "hessian_accumulate leaked the inactive branch");
5249
5250        let ops = t.ops.len();
5251        let mut vals = vec![0.0; ops];
5252        t.forward_into(&x, &mut vals);
5253        let mut col = vec![0.0; 1];
5254        let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
5255        t.hessian_directional(
5256            &vals,
5257            &[1.0],
5258            1.0,
5259            &mut col,
5260            &mut dot,
5261            &mut adj,
5262            &mut adj_dot,
5263        );
5264        assert_eq!(
5265            col[0], 0.0,
5266            "hessian_directional leaked the inactive branch"
5267        );
5268    }
5269
5270    #[test]
5271    fn atan2_grad_and_hessian_match_fd() {
5272        // f = atan2(x0, x1) + x0*x1, away from the origin.
5273        let atan2 = |a: Expr, b: Expr| Expr::Binary(BinOp::Atan2, Box::new(a), Box::new(b));
5274        let e = Expr::Sum(vec![atan2(var(0), var(1)), mul(var(0), var(1))]);
5275        grad_and_hess_match_fd(&e, &[1.2, 0.7], 1e-5);
5276    }
5277
5278    #[test]
5279    fn minmax_grad_and_hessian_match_fd() {
5280        // f = min(x0, x1, x2) + max(x1, x2) + x0*x2
5281        // Point chosen so each list has a UNIQUE strictly-active
5282        // operand, so the subgradient equals the FD slope (the ±h
5283        // probes never cross a kink):
5284        //   min(0.5, 3.0, 2.0) = 0.5  -> active x0
5285        //   max(3.0, 2.0)      = 3.0  -> active x1
5286        let e = Expr::Sum(vec![
5287            Expr::MinList(vec![var(0), var(1), var(2)]),
5288            Expr::MaxList(vec![var(1), var(2)]),
5289            mul(var(0), var(2)),
5290        ]);
5291        grad_and_hess_match_fd(&e, &[0.5, 3.0, 2.0], 1e-5);
5292    }
5293
5294    #[test]
5295    fn minmax_value_and_active_operand() {
5296        // Spot-check the value and that the gradient routes entirely
5297        // through the active operand (zero second derivative).
5298        let e = Expr::Sum(vec![
5299            Expr::MinList(vec![var(0), var(1)]),
5300            Expr::MaxList(vec![var(0), var(1)]),
5301        ]);
5302        let t = Tape::build(&e);
5303        // min(x0,x1) + max(x0,x1) == x0 + x1 for any inputs.
5304        let x = [1.3, -0.4];
5305        assert!((t.eval(&x) - (x[0] + x[1])).abs() < 1e-12);
5306        let mut g = vec![0.0; 2];
5307        t.gradient_seed(&x, 1.0, &mut g);
5308        // min active = x1 (smaller), max active = x0 (larger):
5309        // d/dx0 = 1 (from max), d/dx1 = 1 (from min).
5310        assert!((g[0] - 1.0).abs() < 1e-12, "g0={}", g[0]);
5311        assert!((g[1] - 1.0).abs() < 1e-12, "g1={}", g[1]);
5312    }
5313
5314    #[test]
5315    fn hessian_division_matches_fd() {
5316        // f = x0/x1 + cos(x0)
5317        let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
5318        let t = Tape::build(&e);
5319        fd_check(&t, &[0.5, 1.2], 2, 1e-5);
5320    }
5321
5322    #[test]
5323    fn conditional_value_grad_hessian_active_branch() {
5324        // f = if x0 >= 1 then x0*x1 else x1^2
5325        // The if-then-else differentiates only the active branch; the
5326        // condition (a comparison) contributes no derivative.
5327        let e = cond(
5328            cmp(CmpOp::Ge, var(0), cnst(1.0)),
5329            mul(var(0), var(1)),
5330            pow(var(1), cnst(2.0)),
5331        );
5332        let t = Tape::build(&e);
5333
5334        // x0 = 2 (>= 1) -> "then" branch x0*x1 is active.
5335        let x = [2.0, 5.0];
5336        assert!((t.eval(&x) - 10.0).abs() < 1e-12);
5337        let mut g = vec![0.0; 2];
5338        t.gradient_seed(&x, 1.0, &mut g);
5339        // d(x0*x1) = (x1, x0) = (5, 2)
5340        assert!((g[0] - 5.0).abs() < 1e-10);
5341        assert!((g[1] - 2.0).abs() < 1e-10);
5342        // H[0,1] = 1, diagonals 0. (Stay clear of the x0 = 1 kink.)
5343        fd_check(&t, &x, 2, 1e-5);
5344
5345        // x0 = 0 (< 1) -> "else" branch x1^2 is active; x0 drops out.
5346        let x2 = [0.0, 5.0];
5347        assert!((t.eval(&x2) - 25.0).abs() < 1e-12);
5348        let mut g2 = vec![0.0; 2];
5349        t.gradient_seed(&x2, 1.0, &mut g2);
5350        assert!(g2[0].abs() < 1e-10);
5351        assert!((g2[1] - 10.0).abs() < 1e-10);
5352        fd_check(&t, &x2, 2, 1e-5);
5353    }
5354
5355    #[test]
5356    fn comparison_and_logical_have_zero_derivative() {
5357        // f = (x0 < x1) + (x0 > 0 && x1 > 0) + !(x0 == x1)
5358        // Every term is piecewise-constant in the variables, so the
5359        // gradient must be identically zero away from the kinks.
5360        let lt = cmp(CmpOp::Lt, var(0), var(1));
5361        let and = Expr::And(
5362            Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
5363            Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
5364        );
5365        let notc = Expr::Not(Box::new(cmp(CmpOp::Eq, var(0), var(1))));
5366        let e = add(add(lt, and), notc);
5367        let t = Tape::build(&e);
5368
5369        let x = [1.0, 2.0];
5370        // 1 (1<2) + 1 (both > 0) + 1 (1 != 2) = 3
5371        assert!((t.eval(&x) - 3.0).abs() < 1e-12);
5372        let mut g = vec![0.0; 2];
5373        t.gradient_seed(&x, 1.0, &mut g);
5374        assert!(g[0].abs() < 1e-12, "d/dx0 should be 0, got {}", g[0]);
5375        assert!(g[1].abs() < 1e-12, "d/dx1 should be 0, got {}", g[1]);
5376    }
5377
5378    #[test]
5379    fn logical_or_value() {
5380        // f = (x0 > 0 || x1 > 0)
5381        let e = Expr::Or(
5382            Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
5383            Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
5384        );
5385        let t = Tape::build(&e);
5386        assert!((t.eval(&[-1.0, 3.0]) - 1.0).abs() < 1e-12);
5387        assert!((t.eval(&[-1.0, -3.0]) - 0.0).abs() < 1e-12);
5388    }
5389
5390    /// `hessian_directional` (one forward-over-reverse pass with
5391    /// a seed vector) recovers `H · e_j` for each unit-vector seed,
5392    /// matching column `j` of the dense Hessian computed by
5393    /// `hessian_accumulate`.
5394    fn directional_matches_accumulate(tape: &Tape, x: &[f64], n: usize) {
5395        let vars = tape.variables();
5396        let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5397        let mut pairs = Vec::new();
5398        for (ai, &vi) in vars.iter().enumerate() {
5399            for &vj in &vars[..=ai] {
5400                let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
5401                hess_map.entry((r, c)).or_insert_with(|| {
5402                    let p = pairs.len();
5403                    pairs.push((r, c));
5404                    p
5405                });
5406            }
5407        }
5408        let nnz = pairs.len();
5409        let mut ad = vec![0.0; nnz];
5410        tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
5411
5412        let nops = tape.ops.len();
5413        let mut vals = vec![0.0; nops];
5414        tape.forward_into(x, &mut vals);
5415        let mut dot = vec![0.0; nops];
5416        let mut adj = vec![0.0; nops];
5417        let mut adj_dot = vec![0.0; nops];
5418
5419        for &j in &vars {
5420            let mut seed = vec![0.0; n];
5421            seed[j] = 1.0;
5422            let mut col = vec![0.0; n];
5423            tape.hessian_directional(
5424                &vals,
5425                &seed,
5426                1.0,
5427                &mut col,
5428                &mut dot,
5429                &mut adj,
5430                &mut adj_dot,
5431            );
5432            for &i in &vars {
5433                let (r, c) = if i >= j { (i, j) } else { (j, i) };
5434                let expect = ad[hess_map[&(r, c)]];
5435                assert!(
5436                    (col[i] - expect).abs() < 1e-10,
5437                    "directional H[{i},{j}] = {} vs accumulate {}",
5438                    col[i],
5439                    expect
5440                );
5441            }
5442        }
5443    }
5444
5445    #[test]
5446    fn directional_quadratic_matches_accumulate() {
5447        // f = 3 x0^2 + 2 x0 x1 + x1^2
5448        let e = add(
5449            add(
5450                mul(cnst(3.0), pow(var(0), cnst(2.0))),
5451                mul(mul(cnst(2.0), var(0)), var(1)),
5452            ),
5453            pow(var(1), cnst(2.0)),
5454        );
5455        let t = Tape::build(&e);
5456        directional_matches_accumulate(&t, &[0.5, -0.3], 2);
5457    }
5458
5459    #[test]
5460    fn directional_transcendental_matches_accumulate() {
5461        let e = Expr::Sum(vec![
5462            unary(UnaryOp::Exp, var(0)),
5463            unary(UnaryOp::Sin, var(1)),
5464            unary(UnaryOp::Log, var(0)),
5465            unary(UnaryOp::Sqrt, var(1)),
5466            mul(var(0), var(1)),
5467        ]);
5468        let t = Tape::build(&e);
5469        directional_matches_accumulate(&t, &[1.5, 2.0], 2);
5470    }
5471
5472    #[test]
5473    fn directional_with_division_matches_accumulate() {
5474        let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
5475        let t = Tape::build(&e);
5476        directional_matches_accumulate(&t, &[0.5, 1.2], 2);
5477    }
5478
5479    #[test]
5480    fn hessian_sparsity_separable() {
5481        // f = sin(x0) + x1*x2; couplings: (0,0) from sin, (2,1) from x1*x2
5482        let e = add(unary(UnaryOp::Sin, var(0)), mul(var(1), var(2)));
5483        let t = Tape::build(&e);
5484        let s = t.hessian_sparsity();
5485        assert!(s.contains(&(0, 0)));
5486        assert!(s.contains(&(2, 1)));
5487        assert!(!s.contains(&(1, 0)));
5488        assert!(!s.contains(&(2, 0)));
5489    }
5490
5491    fn count_op<F: Fn(&TapeOp) -> bool>(t: &Tape, pred: F) -> usize {
5492        t.ops.iter().filter(|o| pred(o)).count()
5493    }
5494
5495    #[test]
5496    fn pow_zero_const_folds_to_one() {
5497        // x^0 → 1 (no Pow, no reference to x in the tape)
5498        let e = pow(var(0), cnst(0.0));
5499        let t = Tape::build(&e);
5500        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5501        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Var(_))), 0);
5502        assert!((t.eval(&[7.0]) - 1.0).abs() < 1e-12);
5503    }
5504
5505    #[test]
5506    fn pow_one_passes_through() {
5507        // x^1 → x (no Pow, no Const introduced for the exponent)
5508        let e = pow(var(0), cnst(1.0));
5509        let t = Tape::build(&e);
5510        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5511        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Const(_))), 0);
5512        assert!((t.eval(&[3.5]) - 3.5).abs() < 1e-12);
5513    }
5514
5515    #[test]
5516    fn pow_half_lowers_to_sqrt() {
5517        let e = pow(var(0), cnst(0.5));
5518        let t = Tape::build(&e);
5519        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5520        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Sqrt(_))), 1);
5521        assert!((t.eval(&[16.0]) - 4.0).abs() < 1e-12);
5522    }
5523
5524    #[test]
5525    fn pow_two_lowers_to_single_mul() {
5526        let e = pow(var(0), cnst(2.0));
5527        let t = Tape::build(&e);
5528        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5529        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5530        assert!((t.eval(&[3.0]) - 9.0).abs() < 1e-12);
5531    }
5532
5533    #[test]
5534    fn pow_three_lowers_to_two_muls() {
5535        let e = pow(var(0), cnst(3.0));
5536        let t = Tape::build(&e);
5537        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5538        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 2);
5539        assert!((t.eval(&[2.0]) - 8.0).abs() < 1e-12);
5540    }
5541
5542    #[test]
5543    fn pow_eight_lowers_to_three_muls() {
5544        // Binary expansion: x → x² → x⁴ → x⁸ (3 squarings)
5545        let e = pow(var(0), cnst(8.0));
5546        let t = Tape::build(&e);
5547        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5548        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 3);
5549        assert!((t.eval(&[2.0]) - 256.0).abs() < 1e-12);
5550    }
5551
5552    #[test]
5553    fn pow_negative_two_lowers_to_div() {
5554        // x^-2 → 1 / (x*x)
5555        let e = pow(var(0), cnst(-2.0));
5556        let t = Tape::build(&e);
5557        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5558        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Div(..))), 1);
5559        assert!((t.eval(&[4.0]) - (1.0 / 16.0)).abs() < 1e-12);
5560    }
5561
5562    #[test]
5563    fn pow_large_const_stays_generic() {
5564        // x^9 stays as Pow — beyond the cutoff, generic is cheaper.
5565        let e = pow(var(0), cnst(9.0));
5566        let t = Tape::build(&e);
5567        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5568    }
5569
5570    #[test]
5571    fn pow_non_integer_const_stays_generic() {
5572        // x^1.5 stays as Pow until half-integer handling is added.
5573        let e = pow(var(0), cnst(1.5));
5574        let t = Tape::build(&e);
5575        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5576    }
5577
5578    #[test]
5579    fn pow_const_through_cse_const() {
5580        // Exponent wrapped in Cse — peek_const should still see it.
5581        let two = Arc::new(cnst(2.0));
5582        let e = Expr::Binary(BinOp::Pow, Box::new(var(0)), Box::new(Expr::Cse(two)));
5583        let t = Tape::build(&e);
5584        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5585        assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5586    }
5587
5588    #[test]
5589    fn hessian_pow_three_matches_fd() {
5590        // f = 5 * x0^3 + x0 * x1 — exercises the lowered cubic + cross term.
5591        let e = add(mul(cnst(5.0), pow(var(0), cnst(3.0))), mul(var(0), var(1)));
5592        let t = Tape::build(&e);
5593        fd_check(&t, &[1.7, 0.8], 2, 1e-5);
5594    }
5595
5596    #[test]
5597    fn hessian_pow_negative_matches_fd() {
5598        // f = 1/x0^2 + x1^2 — exercises lowered x^-2 and x^2.
5599        let e = add(pow(var(0), cnst(-2.0)), pow(var(1), cnst(2.0)));
5600        let t = Tape::build(&e);
5601        fd_check(&t, &[1.3, 2.4], 2, 1e-5);
5602    }
5603
5604    #[test]
5605    fn hessian_pow_half_matches_fd() {
5606        // f = sqrt(x0) + x0*x1 (via Pow(_, 0.5) → Sqrt)
5607        let e = add(pow(var(0), cnst(0.5)), mul(var(0), var(1)));
5608        let t = Tape::build(&e);
5609        fd_check(&t, &[2.5, 1.1], 2, 1e-5);
5610    }
5611
5612    #[test]
5613    fn hessian_sparsity_through_cse() {
5614        // body = x0+x1 (CSE). f = body^2 + body.
5615        // d²/dx² of body^2 couples (0,0), (1,0), (1,1).
5616        let body = Arc::new(add(var(0), var(1)));
5617        let e = add(
5618            pow(Expr::Cse(body.clone()), cnst(2.0)),
5619            Expr::Cse(body.clone()),
5620        );
5621        let t = Tape::build(&e);
5622        let s = t.hessian_sparsity();
5623        assert!(s.contains(&(0, 0)));
5624        assert!(s.contains(&(1, 0)));
5625        assert!(s.contains(&(1, 1)));
5626        assert_eq!(s.len(), 3);
5627    }
5628
5629    #[test]
5630    fn pow_forward_tangent_matches_reverse_gradient_at_base_zero() {
5631        // Code review L29: `Pow` first-order tangent disagreed with the
5632        // reverse-mode gradient at base 0. f = x0 ^ x1 keeps a genuine `Pow`
5633        // op (variable exponent is not lowered to a Mul/Sqrt chain). At the
5634        // `.nl` default start x0 = 0, the base derivative d/dx0 (x0^1) = 1 is
5635        // well defined; reverse mode has always computed it, but the forward
5636        // tangent used to guard on `u != 0` and drop it, so Jacobian-vector
5637        // products silently disagreed with the gradient at x = 0. After the
5638        // fix both arms must agree.
5639        let e = pow(var(0), var(1));
5640        let t = Tape::build(&e);
5641        // Guard: the op must survive as a real Pow (not lowered away), else
5642        // this test would no longer exercise the fixed branch.
5643        assert!(
5644            t.ops.iter().any(|op| matches!(op, TapeOp::Pow(_, _))),
5645            "expected a Pow op in the tape; got {:?}",
5646            t.ops
5647        );
5648        let x = [0.0, 1.0];
5649        let n = t.ops.len();
5650
5651        // Reverse-mode gradient w.r.t. x0.
5652        let mut grad = vec![0.0; 2];
5653        t.gradient_seed(&x, 1.0, &mut grad);
5654
5655        // Forward tangent seeded on x0: dot[output] = df/dx0.
5656        let vals = t.forward(&x);
5657        let mut dot = vec![0.0; n];
5658        t.forward_tangent(&vals, 0, &mut dot);
5659        let fwd_dfx0 = dot[n - 1];
5660
5661        assert!(
5662            (grad[0] - 1.0).abs() < 1e-12,
5663            "reverse gradient df/dx0 at base 0 should be 1, got {}",
5664            grad[0]
5665        );
5666        assert!(
5667            (fwd_dfx0 - grad[0]).abs() < 1e-12,
5668            "forward tangent df/dx0 = {fwd_dfx0} must match reverse gradient {} at base 0",
5669            grad[0]
5670        );
5671    }
5672
5673    #[test]
5674    #[should_panic(expected = "external function calls are not supported on the")]
5675    fn hybrid_promoted_cse_with_funcall_reports_clear_message() {
5676        // Code review L34: `HybridTape::build_multi` builds a promoted CSE
5677        // (one shared across ≥2 summands) via `build_recursive` with an empty
5678        // resolver. A funcall inside that promoted body used to panic with the
5679        // misleading "unresolved AMPL funcall id 0" — implying a resolution
5680        // failure — instead of the real reason: funcalls are unsupported on
5681        // the hybrid path. Here the funcall body is shared across two roots so
5682        // it is promoted; assert the clear hybrid-unsupported message fires.
5683        let body = Arc::new(Expr::Funcall {
5684            id: 0,
5685            args: vec![FuncallArg::Real(var(0))],
5686        });
5687        let exprs = vec![
5688            add(Expr::Cse(body.clone()), cnst(1.0)),
5689            add(Expr::Cse(body.clone()), cnst(2.0)),
5690        ];
5691        HybridTape::build_multi(&exprs);
5692    }
5693
5694    /// The single-shooting chain `e_{i+1} = e_i + 0.1*(x_i - e_i*e_i)` as
5695    /// `build_nl_problem` hands it over (gh#1066): each `e_i` one shared
5696    /// `Cse` body, the top level split into two summands so the bodies
5697    /// are promoted. `leaf` is the innermost operand.
5698    fn shooting_chain_roots(k: usize, leaf: Expr) -> Vec<Expr> {
5699        let step = |e: &Expr, i: usize| {
5700            let sq = mul(e.clone(), e.clone());
5701            mul(cnst(0.1), add(var(i), mul(cnst(-1.0), sq)))
5702        };
5703        let mut e = leaf;
5704        for i in 0..k - 1 {
5705            e = Expr::Cse(Arc::new(add(e.clone(), step(&e, i))));
5706        }
5707        let last = step(&e, k - 1);
5708        vec![e, last]
5709    }
5710
5711    /// gh#1066: the promoted-CSE funcall pre-scan walked every path through
5712    /// the shared bodies, `2^K` of them on this chain — K = 20 took 17 s
5713    /// and K = 40 never returned. It is linear in the DAG now, so K = 64
5714    /// builds at once (an exponential walk would need ~`3^64` visits).
5715    #[test]
5716    fn hybrid_build_of_a_shared_square_chain_is_linear() {
5717        let k = 64;
5718        let roots = shooting_chain_roots(k, cnst(0.5));
5719        let tape = HybridTape::build_multi(&roots);
5720        // The value agrees with the flat tape of the whole expression.
5721        let x: Vec<f64> = (0..k).map(|i| 0.01 * i as f64).collect();
5722        let flat = Tape::build(&add(roots[0].clone(), roots[1].clone())).eval(&x);
5723        let mut prelude = vec![0.0; tape.n_prelude_ops()];
5724        tape.forward_prelude(&x, &mut prelude);
5725        let mut local = vec![0.0; tape.max_summand_ops()];
5726        let mut hybrid = 0.0;
5727        for s in &tape.summands {
5728            tape.forward_summand(s, &x, &prelude, &mut local);
5729            hybrid += tape.root_value(s, &local);
5730        }
5731        assert!(
5732            (hybrid - flat).abs() <= 1e-12 * flat.abs().max(1.0),
5733            "{hybrid} vs {flat}"
5734        );
5735    }
5736
5737    /// The linear scan must still see a funcall at the bottom of the same
5738    /// chain: stopping at already-built bodies and at revisits may skip
5739    /// work, never a body that was not checked.
5740    #[test]
5741    #[should_panic(expected = "external function calls are not supported on the")]
5742    fn hybrid_funcall_under_a_shared_square_chain_still_reports() {
5743        let leaf = Expr::Funcall {
5744            id: 0,
5745            args: vec![FuncallArg::Real(var(0))],
5746        };
5747        HybridTape::build_multi(&shooting_chain_roots(64, leaf));
5748    }
5749
5750    /// The directional hybrid Hessian (issue #557) against the flat
5751    /// [`Tape::hessian_directional`] it replaces: three summands sharing one
5752    /// transcendental CSE body, each with its own multiplier, swept with a
5753    /// dense two-variable seed. The per-summand *local* contributions are the
5754    /// exact arithmetic of the flat tape; the shared-prelude contribution
5755    /// folds the multipliers into the boundary adjoints and runs one unit-
5756    /// weight prelude sweep, which reassociates the floating-point products —
5757    /// mathematically identical by linearity of reverse-over-tangent in the
5758    /// adjoint seeds, so the comparison is at near-machine tolerance rather
5759    /// than bitwise.
5760    #[test]
5761    fn directional_hybrid_hessian_matches_flat_directional() {
5762        // body = exp(0.5*(x0 + x1)), shared by all three summands:
5763        //   s0 = body^2,  s1 = body * x2,  s2 = sin(body) + x2^2.
5764        let body = Arc::new(unary(UnaryOp::Exp, mul(cnst(0.5), add(var(0), var(1)))));
5765        let exprs = vec![
5766            pow(Expr::Cse(body.clone()), cnst(2.0)),
5767            mul(Expr::Cse(body.clone()), var(2)),
5768            add(
5769                unary(UnaryOp::Sin, Expr::Cse(body.clone())),
5770                pow(var(2), cnst(2.0)),
5771            ),
5772        ];
5773        let weights = [1.25, -0.75, 2.5];
5774        let x = [0.3, -0.1, 0.7];
5775        let n = 3;
5776
5777        let hybrid = HybridTape::build_multi(&exprs);
5778        assert!(
5779            hybrid.n_prelude_ops() > 0,
5780            "a CSE shared by 3 roots must be promoted into the prelude"
5781        );
5782
5783        // Flat reference: independent tapes, one directional pass each.
5784        let seeds = [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]];
5785        for seed in &seeds {
5786            let mut flat_out = vec![0.0; n];
5787            for (e, &wt) in exprs.iter().zip(&weights) {
5788                let t = Tape::build(e);
5789                let vals = t.forward(&x);
5790                let m = t.ops.len();
5791                let (mut dot, mut adj, mut adj_dot) = (vec![0.0; m], vec![0.0; m], vec![0.0; m]);
5792                t.hessian_directional(
5793                    &vals,
5794                    seed,
5795                    wt,
5796                    &mut flat_out,
5797                    &mut dot,
5798                    &mut adj,
5799                    &mut adj_dot,
5800                );
5801            }
5802
5803            let mut hyb_out = vec![0.0; n];
5804            let np = hybrid.n_prelude_ops();
5805            let ml = hybrid.max_summand_ops();
5806            let mut prelude_vals = vec![0.0; np];
5807            let mut prelude_dot = vec![0.0; np];
5808            let mut prelude_adj = vec![0.0; np];
5809            let mut prelude_adj_dot = vec![0.0; np];
5810            let (mut local_dot, mut local_adj, mut local_adj_dot) =
5811                (vec![0.0; ml], vec![0.0; ml], vec![0.0; ml]);
5812            // The per-color prelude reach the caller is required to pass:
5813            // the union of the color's summands' `prelude_reach`, ascending.
5814            // Built here rather than passing `0..np` so the test exercises the
5815            // narrowed walk `eval_h` actually performs — a reach that wrongly
5816            // omitted a slot would leave its tangent stale and show up below.
5817            let creach: Vec<u32> = {
5818                let mut u: BTreeSet<u32> = BTreeSet::new();
5819                for s in &hybrid.summands {
5820                    u.extend(s.prelude_reach.iter().map(|&p| p as u32));
5821                }
5822                u.into_iter().collect()
5823            };
5824            hybrid.forward_prelude(&x, &mut prelude_vals);
5825            hybrid.prelude_tangent(&prelude_vals, seed, &creach, &mut prelude_dot);
5826            for (s, &wt) in hybrid.summands.iter().zip(&weights) {
5827                let mut local_vals = vec![0.0; s.ops.len()];
5828                hybrid.forward_summand(s, &x, &prelude_vals, &mut local_vals);
5829                hybrid.hessian_summand_directional(
5830                    s,
5831                    &local_vals,
5832                    &prelude_dot,
5833                    seed,
5834                    wt,
5835                    &mut hyb_out,
5836                    &mut local_dot,
5837                    &mut local_adj,
5838                    &mut local_adj_dot,
5839                    &mut prelude_adj,
5840                    &mut prelude_adj_dot,
5841                );
5842            }
5843            hybrid.prelude_reverse_directional(
5844                &prelude_vals,
5845                &prelude_dot,
5846                &creach,
5847                &mut hyb_out,
5848                &mut prelude_adj,
5849                &mut prelude_adj_dot,
5850            );
5851
5852            for k in 0..n {
5853                let scale = flat_out[k].abs().max(1.0);
5854                assert!(
5855                    (hyb_out[k] - flat_out[k]).abs() <= 1e-13 * scale,
5856                    "seed {seed:?} entry {k}: hybrid {} vs flat {}",
5857                    hyb_out[k],
5858                    flat_out[k]
5859                );
5860            }
5861            assert!(
5862                hyb_out.iter().any(|v| *v != 0.0),
5863                "an all-zero H·s is no comparison"
5864            );
5865            // The accumulators must come back all-zero, ready for the
5866            // next color (`prelude_reverse_directional`'s contract).
5867            assert!(prelude_adj.iter().all(|v| *v == 0.0));
5868            assert!(prelude_adj_dot.iter().all(|v| *v == 0.0));
5869        }
5870    }
5871}