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