Skip to main content

gam_math/
jet_scalar.rs

1//! Order-specific Taylor-jet SCALAR algebras (#932 cutover, doc §A).
2//!
3//! [`crate::jet_tower::Tower4`] carries the full value/gradient/Hessian/`t3`/`t4`
4//! tensor stack: it answers EVERY channel a [`super::row_kernel::RowKernel`]
5//! consumer can ask for, but at `K = 9` that is a ~50 KiB per-row object whose
6//! by-value copies overflowed the stack and timed out the location-scale fit.
7//! The cutover therefore does NOT instantiate the dense `Tower4<9>` per row; it
8//! carries, per consumer, only the one channel that consumer serves:
9//!
10//! | consumer | channel | scalar here | K=9 size |
11//! |---|---|---|---|
12//! | inner Newton / `row_kernel` | `(v, g, H)` | [`Order2`] | 728 B |
13//! | `row_third_contracted(dir)` | `Σ_c ℓ_{abc} dir_c` | [`OneSeed`] | 1.46 KiB |
14//! | `row_fourth_contracted(u, v)` | `Σ_{cd} ℓ_{abcd} u_c v_d` | [`TwoSeed`] | 2.8 KiB |
15//!
16//! Each is built on [`Order2`] (value/grad/Hessian), which is the production
17//! [`crate::jet_tower::Tower2`] re-expressed behind a generic interface: a row
18//! loss written ONCE against [`JetScalar`] re-instantiates at whatever order /
19//! representation a consumer needs, with the contraction folded INTO the
20//! differentiation (the nilpotent ε / δ directions), so `t3` / `t4` are never
21//! materialised. The single source of truth is the same one expression — the
22//! genus of #736 cross-block drift cannot reappear because there is no separate
23//! channel to forget.
24//!
25//! # Why each scalar is exact (doc §A.1–A.3)
26//!
27//! * [`Order2`] is the order-≤2 truncation of the Leibniz / Faà di Bruno rules.
28//!   Those order-2 terms read ONLY the order-≤2 channels of their inputs (see
29//!   [`crate::jet_tower::Tower4::mul`]: `out.h[i][j]` never touches `t3`/`t4`),
30//!   so its `(v, g, H)` is BIT-IDENTICAL to a full `Tower4<K>` — and identical
31//!   to [`crate::jet_tower::Tower2`], over which it is a thin newtype.
32//! * [`OneSeed`] carries an [`Order2`] base plus one nilpotent ε (`ε² = 0`)
33//!   holding another [`Order2`]. Seeding ε with the fixed direction `u` makes the
34//!   ε-component of the Hessian channel the contracted third `Σ_c ℓ_{abc} u_c`
35//!   (the nilpotent implements `d/dτ|₀` of `ℓ_{ab}(p + τu)` exactly).
36//! * [`TwoSeed`] carries an [`Order2`] base plus ε, δ (`ε² = δ² = 0`, `εδ`
37//!   retained) — four [`Order2`] parts. Seeding ε, δ with `u, v` makes the
38//!   εδ-component of the Hessian channel the contracted fourth
39//!   `Σ_{cd} ℓ_{abcd} u_c v_d` (the single mixed `∂_σ∂_ρ|₀` term, no `σ²`/`ρ²`
40//!   contamination).
41//!
42//! # Stability discipline
43//!
44//! As in [`crate::jet_tower`], humans own primitive stability and the algebra
45//! owns combinatorics: tail-critical special functions enter ONLY as
46//! hand-certified `[f64; 5]` derivative stacks through [`crate::nested_dual::JetField::compose_unary`]
47//! (each scalar consumes the leading entries its order needs), never by
48//! differentiating an unstable primal.
49//!
50//! # Production scalars and the test-only all-channels oracle
51//!
52//! The `JetScalar` trait below is production: it is the bound on the canonical
53//! [`crate::jet_tower::RowProgram::eval`] seam a family row loss is written
54//! against. The order-specific scalars that *consume* it —
55//! [`Order2`] (value/grad/Hessian), [`OneSeed`] (contracted third) and
56//! [`TwoSeed`] (contracted fourth) — are production: the survival location-scale
57//! `RowKernel<9>` builds its joint Hessian / directional derivatives through them
58//! (`survival::location_scale::row_kernel`), paying only the small packed scalar
59//! per row instead of the ~50 KiB dense [`crate::jet_tower::Tower4`].
60//!
61//! The [`crate::jet_tower::Tower4`] all-channels `JetScalar` impl is test-only: it
62//! is the oracle that pins the contracted scalars against the dense
63//! value/grad/Hessian/`t3`/`t4` truth, so it lives in the `#[cfg(test)]` module.
64
65/// Symmetric coefficient operator used by the universal quadratic-form jet
66/// primitive.
67///
68/// Implementations own representation-specific matrix access (diagonal,
69/// dense, low rank, structured, or matrix free). Jet scalars own the one
70/// derivative rule for `x' A x`; callers never spell its chain rule.
71pub trait SymmetricQuadraticCoefficients {
72    /// Input-space dimension of the symmetric operator.
73    fn dimension(&self) -> usize;
74
75    /// Compute `output = A * input` without materializing `A`.
76    fn multiply(&self, input: &[f64], output: &mut [f64]);
77
78    /// One symmetric coefficient `A[row, column]`.
79    fn coefficient(&self, row: usize, column: usize) -> f64;
80
81    /// Visit the operator's input-space upper triangle without materializing a
82    /// dense matrix. The default probes one basis direction at a time through
83    /// [`Self::multiply`], preserving matrix-free operators. Representations
84    /// with direct structure should override this hook so a packed curvature
85    /// lowering is O(K), O(K²), or O(K²R) as appropriate instead of performing
86    /// K full operator applications.
87    ///
88    /// `direction` and `projected` are caller-owned scratch of exact length
89    /// [`Self::dimension`]. Every `(row, column)` with `row <= column` is
90    /// visited exactly once, including structural zeros.
91    fn visit_upper_triangle(
92        &self,
93        direction: &mut [f64],
94        projected: &mut [f64],
95        mut visit: impl FnMut(usize, usize, f64),
96    ) {
97        let dimension = self.dimension();
98        assert_eq!(direction.len(), dimension);
99        assert_eq!(projected.len(), dimension);
100        direction.fill(0.0);
101        for column in 0..dimension {
102            direction[column] = 1.0;
103            self.multiply(direction, projected);
104            direction[column] = 0.0;
105            for row in 0..=column {
106                visit(row, column, projected[row]);
107            }
108        }
109    }
110
111    /// Evaluate the primal quadratic form without materializing the input as a
112    /// separate `f64` vector. Structured operators should override this to
113    /// preserve their representation's natural complexity (for example O(KR)
114    /// for a K-by-R low-rank factor).
115    fn quadratic_value<T, F>(&self, inputs: &[T], value: F) -> f64
116    where
117        F: Fn(&T) -> f64,
118    {
119        assert_eq!(
120            inputs.len(),
121            self.dimension(),
122            "symmetric quadratic-form dimension mismatch"
123        );
124        let mut out = 0.0;
125        for row in 0..inputs.len() {
126            let row_value = value(&inputs[row]);
127            out += self.coefficient(row, row) * row_value * row_value;
128            for column in row + 1..inputs.len() {
129                out += 2.0 * self.coefficient(row, column) * row_value * value(&inputs[column]);
130            }
131        }
132        out
133    }
134}
135
136fn symmetric_quadratic_form_default<T, C>(
137    inputs: &[T],
138    coefficients: &C,
139    constant: impl Fn(f64) -> T,
140    add: impl Fn(&T, &T) -> T,
141    mul: impl Fn(&T, &T) -> T,
142    scale: impl Fn(&T, f64) -> T,
143) -> T
144where
145    C: SymmetricQuadraticCoefficients,
146{
147    assert_eq!(
148        inputs.len(),
149        coefficients.dimension(),
150        "symmetric quadratic-form dimension mismatch"
151    );
152    let mut out = constant(0.0);
153    for row in 0..inputs.len() {
154        let diagonal = mul(&inputs[row], &inputs[row]);
155        out = add(&out, &scale(&diagonal, coefficients.coefficient(row, row)));
156        for column in row + 1..inputs.len() {
157            let cross = mul(&inputs[row], &inputs[column]);
158            out = add(
159                &out,
160                &scale(&cross, 2.0 * coefficients.coefficient(row, column)),
161            );
162        }
163    }
164    out
165}
166
167fn linear_combination_default<T>(
168    inputs: &[T],
169    weights: &[f64],
170    constant: impl Fn(f64) -> T,
171    add: impl Fn(&T, &T) -> T,
172    scale: impl Fn(&T, f64) -> T,
173) -> T {
174    assert_eq!(
175        inputs.len(),
176        weights.len(),
177        "linear-combination dimension mismatch"
178    );
179    inputs
180        .iter()
181        .zip(weights)
182        .fold(constant(0.0), |sum, (input, &weight)| {
183            add(&sum, &scale(input, weight))
184        })
185}
186
187fn multiply_add_default<T>(
188    left: &T,
189    right: &T,
190    addend: &T,
191    mul: impl Fn(&T, &T) -> T,
192    add: impl Fn(&T, &T) -> T,
193) -> T {
194    add(&mul(left, right), addend)
195}
196
197fn composed_sum_default<T>(
198    inputs: &[T],
199    derivative_stacks: &[[f64; 5]],
200    constant: impl Fn(f64) -> T,
201    add: impl Fn(&T, &T) -> T,
202    compose: impl Fn(&T, [f64; 5]) -> T,
203) -> T {
204    assert_eq!(
205        inputs.len(),
206        derivative_stacks.len(),
207        "composed-sum term-count mismatch"
208    );
209    inputs
210        .iter()
211        .zip(derivative_stacks)
212        .fold(constant(0.0), |sum, (input, &stack)| {
213            add(&sum, &compose(input, stack))
214        })
215}
216
217fn affine_compose_default<T>(
218    input: &T,
219    input_scale: f64,
220    input_shift: f64,
221    derivative_stack: [f64; 5],
222    scale: impl Fn(&T, f64) -> T,
223    add_constant: impl Fn(&T, f64) -> T,
224    compose: impl Fn(&T, [f64; 5]) -> T,
225) -> T {
226    compose(
227        &add_constant(&scale(input, input_scale), input_shift),
228        derivative_stack,
229    )
230}
231
232fn affine_composed_sum_default<T>(
233    inputs: &[T],
234    input_scales: &[f64],
235    derivative_stacks: &[[f64; 5]],
236    constant: impl Fn(f64) -> T,
237    add: impl Fn(&T, &T) -> T,
238    scale: impl Fn(&T, f64) -> T,
239    add_constant: impl Fn(&T, f64) -> T,
240    compose: impl Fn(&T, [f64; 5]) -> T,
241) -> T {
242    assert_eq!(inputs.len(), input_scales.len());
243    assert_eq!(inputs.len(), derivative_stacks.len());
244    inputs.iter().zip(input_scales).zip(derivative_stacks).fold(
245        constant(0.0),
246        |sum, ((input, &input_scale), &stack)| {
247            add(
248                &sum,
249                &affine_compose_default(
250                    input,
251                    input_scale,
252                    0.0,
253                    stack,
254                    &scale,
255                    &add_constant,
256                    &compose,
257                ),
258            )
259        },
260    )
261}
262
263fn shared_multiply_add_affine_composed_sum_default<T, const N: usize>(
264    lefts: &[&T; N],
265    right: &T,
266    addend: &T,
267    addend_scales: &[f64; N],
268    input_scales: &[f64; N],
269    derivative_stacks: &[[f64; 5]; N],
270    constant: impl Fn(f64) -> T,
271    add: impl Fn(&T, &T) -> T,
272    mul: impl Fn(&T, &T) -> T,
273    scale: impl Fn(&T, f64) -> T,
274    multiply_add: impl Fn(&T, &T, &T) -> T,
275    affine_compose: impl Fn(&T, f64, f64, [f64; 5]) -> T,
276) -> T {
277    let (representatives, term_sources, source_count) =
278        canonical_shared_source_schedule(|term, representative| {
279            std::ptr::eq(lefts[term], lefts[representative])
280                && addend_scales[term] == addend_scales[representative]
281        });
282    let (value, source_derivatives) =
283        aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
284    (0..source_count).fold(constant(value), |sum, source| {
285        let term = representatives[source];
286        let inner = if addend_scales[term] == 0.0 {
287            mul(lefts[term], right)
288        } else if addend_scales[term] == 1.0 {
289            multiply_add(lefts[term], right, addend)
290        } else {
291            multiply_add(lefts[term], right, &scale(addend, addend_scales[term]))
292        };
293        let composed = affine_compose(&inner, 1.0, 0.0, source_derivatives[source]);
294        add(&sum, &composed)
295    })
296}
297
298/// Canonical representatives for repeated inner sources in a shared composed sum.
299///
300/// Backends define source identity at their representation boundary: eager jets
301/// use reference identity, while compiled graphs use node identity. Terms are
302/// equivalent only when their left operand and structural addend coefficient
303/// name the same inner expression. The returned term-to-source map lets every
304/// backend sum outer first/second coefficients before propagating that source.
305#[inline(always)]
306pub(crate) fn canonical_shared_source_schedule<const N: usize>(
307    mut equivalent: impl FnMut(usize, usize) -> bool,
308) -> ([usize; N], [usize; N], usize) {
309    let mut representatives = [0; N];
310    let mut term_sources = [0; N];
311    let mut source_count = 0;
312    for term in 0..N {
313        let mut source = 0;
314        while source < source_count && !equivalent(term, representatives[source]) {
315            source += 1;
316        }
317        if source == source_count {
318            representatives[source] = term;
319            source_count += 1;
320        }
321        term_sources[term] = source;
322    }
323    (representatives, term_sources, source_count)
324}
325
326/// Sum each source's affine-scaled derivative stack through fourth order.
327///
328/// For a term `f(s*x)`, derivative order `r` contributes `f^(r) * s^r`.
329/// The primal value is accumulated once in original term order, while source
330/// stacks deliberately keep a zero value channel so adding their composed jets
331/// cannot perturb that authoritative primal sum.
332#[inline(always)]
333pub(crate) fn aggregate_shared_source_derivatives<const N: usize>(
334    term_sources: &[usize; N],
335    input_scales: &[f64; N],
336    derivative_stacks: &[[f64; 5]; N],
337) -> (f64, [[f64; 5]; N]) {
338    let mut value = 0.0;
339    let mut source_derivatives = [[0.0; 5]; N];
340    for term in 0..N {
341        value += derivative_stacks[term][0];
342        let source = term_sources[term];
343        let mut scale_power = input_scales[term];
344        for order in 1..5 {
345            source_derivatives[source][order] += derivative_stacks[term][order] * scale_power;
346            scale_power *= input_scales[term];
347        }
348    }
349    (value, source_derivatives)
350}
351
352/// A truncated-Taylor scalar carrying derivatives in `K` primaries.
353///
354/// All concrete scalars here ([`Order2`], [`OneSeed`], [`TwoSeed`]) and the full
355/// [`crate::jet_tower::Tower4`] implement the SAME algebra; only the carried
356/// channel set differs. A row loss written once against this interface yields a
357/// different channel set per instantiation, all exact for the channel they serve
358/// (doc §A.0).
359pub trait JetScalar<const K: usize>: crate::nested_dual::JetField + Copy {
360    /// A constant: value `c`, every derivative channel zero.
361    fn constant(c: f64) -> Self;
362
363    /// The seeded variable `p_axis` at value `x`: unit first derivative in slot
364    /// `axis`, all higher channels zero. (The nilpotent / cross channels of the
365    /// directional scalars are seeded zero — callers set ε/δ directions through
366    /// the scalar-specific [`OneSeed::seed_direction`] / [`TwoSeed::seed`].)
367    fn variable(x: f64, axis: usize) -> Self;
368
369    /// Evaluate `inputs' A inputs` from one universal semantic primitive.
370    /// Order-specific scalars may lower the mechanically derived channels
371    /// directly; the default is the exact scalar program over `mul/add/scale`.
372    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
373        inputs: &[Self],
374        coefficients: &C,
375    ) -> Self {
376        symmetric_quadratic_form_default(
377            inputs,
378            coefficients,
379            Self::constant,
380            crate::nested_dual::JetField::add,
381            crate::nested_dual::JetField::mul,
382            crate::nested_dual::JetField::scale,
383        )
384    }
385
386    /// Evaluate `sum_i weights[i] * inputs[i]` in one semantic primitive.
387    fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
388        linear_combination_default(
389            inputs,
390            weights,
391            Self::constant,
392            crate::nested_dual::JetField::add,
393            crate::nested_dual::JetField::scale,
394        )
395    }
396
397    /// Add a primal constant without changing derivative channels.
398    fn add_constant(&self, constant: f64) -> Self {
399        self.add(&Self::constant(constant))
400    }
401
402    /// Evaluate `self * right + addend` in one semantic primitive.
403    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
404        multiply_add_default(
405            self,
406            right,
407            addend,
408            crate::nested_dual::JetField::mul,
409            crate::nested_dual::JetField::add,
410        )
411    }
412
413    /// Sum unary compositions directly from certified derivative stacks.
414    fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
415        composed_sum_default(
416            inputs,
417            derivative_stacks,
418            Self::constant,
419            crate::nested_dual::JetField::add,
420            crate::nested_dual::JetField::compose_unary,
421        )
422    }
423
424    /// Exact product as an explicit compiled graph node.
425    fn product(&self, right: &Self) -> Self {
426        self.mul(right)
427    }
428
429    /// Compose a certified outer stack after the affine map
430    /// `u = input_scale * self + input_shift`.
431    fn affine_compose(
432        &self,
433        input_scale: f64,
434        input_shift: f64,
435        derivative_stack: [f64; 5],
436    ) -> Self {
437        affine_compose_default(
438            self,
439            input_scale,
440            input_shift,
441            derivative_stack,
442            crate::nested_dual::JetField::scale,
443            Self::add_constant,
444            crate::nested_dual::JetField::compose_unary,
445        )
446    }
447
448    /// Sum unary compositions whose inputs each carry an affine scale.
449    fn affine_composed_sum(
450        inputs: &[Self],
451        input_scales: &[f64],
452        derivative_stacks: &[[f64; 5]],
453    ) -> Self {
454        affine_composed_sum_default(
455            inputs,
456            input_scales,
457            derivative_stacks,
458            Self::constant,
459            crate::nested_dual::JetField::add,
460            crate::nested_dual::JetField::scale,
461            Self::add_constant,
462            crate::nested_dual::JetField::compose_unary,
463        )
464    }
465
466    /// Evaluate
467    /// `Σ_i f_i(input_scale_i · (left_i · right + addend_scale_i · addend))`
468    /// from the certified derivative stack of each `f_i`. The shared operands
469    /// make expression-level common subexpressions explicit, so optimized
470    /// backends apply their inherited derivative channels once. Expression
471    /// arity is part of the type, and borrowed operands never copy a full tower.
472    /// An exact-zero addend scale (either sign of IEEE zero) removes that addend
473    /// from the corresponding term entirely.
474    fn shared_multiply_add_affine_composed_sum<const N: usize>(
475        lefts: &[&Self; N],
476        right: &Self,
477        addend: &Self,
478        addend_scales: &[f64; N],
479        input_scales: &[f64; N],
480        derivative_stacks: &[[f64; 5]; N],
481    ) -> Self {
482        shared_multiply_add_affine_composed_sum_default(
483            lefts,
484            right,
485            addend,
486            addend_scales,
487            input_scales,
488            derivative_stacks,
489            Self::constant,
490            crate::nested_dual::JetField::add,
491            crate::nested_dual::JetField::mul,
492            crate::nested_dual::JetField::scale,
493            Self::multiply_add,
494            Self::affine_compose,
495        )
496    }
497
498    // The scalar-field algebra — `value`, `add`, `sub`, `mul`, `neg`, `scale`,
499    // and the single Faà di Bruno `compose_unary([f64; 5])` — is declared ONCE on
500    // the shared [`crate::nested_dual::JetField`] supertrait (which the runtime-`p`
501    // flex jets also extend), so it is not re-declared here. The value/composition
502    // convention is identical (`u = self.value()`, the `[f64; 5]` stack shape
503    // [`crate::jet_tower::Tower4`] consumes). The helpers below build on it.
504
505    /// Compose with a unary special-function whose derivative stack is built from
506    /// the scalar base value through `stack_fn`. This evaluates
507    /// `stack_fn(self.value())` once and forwards to
508    /// [`compose_unary`](Self::compose_unary), so it is bit-identical to the
509    /// explicit `self.compose_unary(stack_fn(self.value()))` form.
510    fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
511        self.compose_unary(stack_fn(self.value()))
512    }
513
514    /// `e^self`. Convenience for tame arguments (see module stability note).
515    fn exp(&self) -> Self {
516        let e = self.value().exp();
517        self.compose_unary([e, e, e, e, e])
518    }
519
520    /// `√self`. Caller guarantees positivity.
521    fn sqrt(&self) -> Self {
522        let u = self.value();
523        let s = u.sqrt();
524        self.compose_unary([
525            s,
526            0.5 / s,
527            -0.25 / (u * s),
528            0.375 / (u * u * s),
529            -0.9375 / (u * u * u * s),
530        ])
531    }
532
533    /// `ln(self)`. Caller guarantees positivity. Same derivative stack
534    /// [`crate::jet_tower::Tower4::ln`] uses, so any program written over both
535    /// matches term-for-term.
536    fn ln(&self) -> Self {
537        let u = self.value();
538        let r = 1.0 / u;
539        self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
540    }
541
542    /// `1/self`.
543    fn recip(&self) -> Self {
544        let r = 1.0 / self.value();
545        let r2 = r * r;
546        self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
547    }
548
549    /// `self^a` for real exponent `a`. Caller guarantees a positive base.
550    /// Mirrors [`crate::jet_tower::Tower4::powf`] (falling-factorial stack).
551    fn powf(&self, a: f64) -> Self {
552        let u = self.value();
553        self.compose_unary([
554            u.powf(a),
555            a * u.powf(a - 1.0),
556            a * (a - 1.0) * u.powf(a - 2.0),
557            a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
558            a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
559        ])
560    }
561
562    /// `ln Γ(self)`. Caller guarantees a positive argument. Uses the SAME
563    /// hand-certified derivative stack [`crate::jet_tower::Tower4::ln_gamma`]
564    /// consumes ([`crate::jet_tower::ln_gamma_derivative_stack`]), so any
565    /// program written over both matches term-for-term.
566    fn ln_gamma(&self) -> Self {
567        self.compose_unary(crate::jet_tower::ln_gamma_derivative_stack(self.value()))
568    }
569
570    /// `ψ(self) = d/dx ln Γ(x)` (digamma). Caller guarantees a positive
571    /// argument. Same hand-certified stack
572    /// [`crate::jet_tower::digamma_derivative_stack`].
573    fn digamma(&self) -> Self {
574        self.compose_unary(crate::jet_tower::digamma_derivative_stack(self.value()))
575    }
576}
577
578/// Lift a fixed-width primary jet through an independent second-order
579/// direction. Primary variables live in the inner scalar; constructing one
580/// must not seed the outer direction, which is reserved for a family or
581/// hyperparameter derivative selected by the caller.
582impl<S, const K: usize> JetScalar<K> for crate::nested_dual::Dual2<S>
583where
584    S: JetScalar<K>,
585{
586    #[inline]
587    fn constant(c: f64) -> Self {
588        Self {
589            v: S::constant(c),
590            g: S::constant(0.0),
591            h: S::constant(0.0),
592        }
593    }
594
595    #[inline]
596    fn variable(x: f64, axis: usize) -> Self {
597        Self {
598            v: S::variable(x, axis),
599            g: S::constant(0.0),
600            h: S::constant(0.0),
601        }
602    }
603}
604
605/// A Taylor-jet scalar whose primary dimension is selected at runtime.
606///
607/// This is the dimensioned counterpart of [`JetScalar`].  Its algebra is the
608/// same; only the constructors receive the row's actual primary count.  The
609/// fixed-size scalar implementations below bridge to this trait as well, which
610/// lets one row program serve both the const-generic derivative oracles and the
611/// runtime-sized production backends without duplicating the expression.
612pub trait RuntimeJetScalar<'arena>: Clone {
613    /// Storage arena used by runtime-backed scalars. Fixed derivative oracles
614    /// use the unit type because their storage is inline.
615    type Workspace: ?Sized;
616
617    /// A constant in a `dimension`-primary algebra.
618    fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
619    /// A seeded variable in a `dimension`-primary algebra.
620    fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
621
622    /// Evaluate `inputs' A inputs` from the same universal semantic primitive
623    /// as [`JetScalar::symmetric_quadratic_form`].
624    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
625        inputs: &[Self],
626        coefficients: &C,
627        dimension: usize,
628        workspace: &'arena Self::Workspace,
629    ) -> Self {
630        symmetric_quadratic_form_default(
631            inputs,
632            coefficients,
633            |value| Self::constant(value, dimension, workspace),
634            Self::add,
635            Self::mul,
636            Self::scale,
637        )
638    }
639
640    /// Evaluate `sum_i weights[i] * inputs[i]` in one semantic primitive.
641    fn linear_combination(
642        inputs: &[Self],
643        weights: &[f64],
644        dimension: usize,
645        workspace: &'arena Self::Workspace,
646    ) -> Self {
647        linear_combination_default(
648            inputs,
649            weights,
650            |value| Self::constant(value, dimension, workspace),
651            Self::add,
652            Self::scale,
653        )
654    }
655
656    /// Add a primal constant without changing derivative channels.
657    fn add_constant(&self, constant: f64, workspace: &'arena Self::Workspace) -> Self {
658        self.add(&Self::constant(constant, self.dimension(), workspace))
659    }
660
661    /// Evaluate `self * right + addend` in one semantic primitive.
662    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
663        multiply_add_default(self, right, addend, Self::mul, Self::add)
664    }
665
666    /// Sum unary compositions directly from certified derivative stacks.
667    fn composed_sum(
668        inputs: &[Self],
669        derivative_stacks: &[[f64; 5]],
670        dimension: usize,
671        workspace: &'arena Self::Workspace,
672    ) -> Self {
673        composed_sum_default(
674            inputs,
675            derivative_stacks,
676            |value| Self::constant(value, dimension, workspace),
677            Self::add,
678            Self::compose_unary,
679        )
680    }
681
682    /// Exact product as an explicit compiled graph node.
683    fn product(&self, right: &Self) -> Self {
684        self.mul(right)
685    }
686
687    /// Compose a certified outer stack after an affine input map.
688    fn affine_compose(
689        &self,
690        input_scale: f64,
691        input_shift: f64,
692        derivative_stack: [f64; 5],
693        workspace: &'arena Self::Workspace,
694    ) -> Self {
695        affine_compose_default(
696            self,
697            input_scale,
698            input_shift,
699            derivative_stack,
700            Self::scale,
701            |value, constant| value.add_constant(constant, workspace),
702            Self::compose_unary,
703        )
704    }
705
706    /// Sum unary compositions whose inputs each carry an affine scale.
707    fn affine_composed_sum(
708        inputs: &[Self],
709        input_scales: &[f64],
710        derivative_stacks: &[[f64; 5]],
711        dimension: usize,
712        workspace: &'arena Self::Workspace,
713    ) -> Self {
714        affine_composed_sum_default(
715            inputs,
716            input_scales,
717            derivative_stacks,
718            |value| Self::constant(value, dimension, workspace),
719            Self::add,
720            Self::scale,
721            |value, constant| value.add_constant(constant, workspace),
722            Self::compose_unary,
723        )
724    }
725
726    /// Runtime-dimension lowering of
727    /// `Σ_i f_i(input_scale_i · (left_i · right + addend_scale_i · addend))`
728    /// from certified derivative stacks. Const expression arity replaces an
729    /// implementation-specific term cap, while shared operands expose universal
730    /// common-subexpression elimination. If every addend scale is exact zero
731    /// (including `-0.0`), the addend has no dimension, workspace, or
732    /// derivative-channel obligations.
733    fn shared_multiply_add_affine_composed_sum<const N: usize>(
734        lefts: &[&Self; N],
735        right: &Self,
736        addend: &Self,
737        addend_scales: &[f64; N],
738        input_scales: &[f64; N],
739        derivative_stacks: &[[f64; 5]; N],
740        dimension: usize,
741        workspace: &'arena Self::Workspace,
742    ) -> Self {
743        shared_multiply_add_affine_composed_sum_default(
744            lefts,
745            right,
746            addend,
747            addend_scales,
748            input_scales,
749            derivative_stacks,
750            |value| Self::constant(value, dimension, workspace),
751            Self::add,
752            Self::mul,
753            Self::scale,
754            Self::multiply_add,
755            |input, scale, shift, stack| input.affine_compose(scale, shift, stack, workspace),
756        )
757    }
758    /// Number of primary derivative axes carried by this scalar.
759    fn dimension(&self) -> usize;
760    /// Value channel.
761    fn value(&self) -> f64;
762    /// Exact truncated sum.
763    fn add(&self, o: &Self) -> Self;
764    /// Exact truncated difference.
765    fn sub(&self, o: &Self) -> Self;
766    /// Exact truncated product.
767    fn mul(&self, o: &Self) -> Self;
768    /// Negate every channel.
769    fn neg(&self) -> Self;
770    /// Scale every channel.
771    fn scale(&self, s: f64) -> Self;
772    /// Exact unary composition from the certified derivative stack.
773    fn compose_unary(&self, d: [f64; 5]) -> Self;
774
775    /// `e^self`.
776    fn exp(&self) -> Self {
777        let e = self.value().exp();
778        self.compose_unary([e, e, e, e, e])
779    }
780
781    /// `ln(self)`. Caller guarantees positivity. This is the runtime-width
782    /// counterpart of [`JetScalar::ln`] and uses the identical certified
783    /// derivative stack.
784    fn ln(&self) -> Self {
785        let u = self.value();
786        let r = 1.0 / u;
787        self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
788    }
789
790    /// `1/self`.
791    fn recip(&self) -> Self {
792        let r = 1.0 / self.value();
793        let r2 = r * r;
794        self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
795    }
796}
797
798/// Zero-order runtime scalar for evaluating the primal value of a canonical
799/// runtime-width row program without constructing derivative channels.
800///
801/// The primary dimension remains part of the scalar so it obeys the same
802/// runtime algebra contract as derivative-carrying implementations. Unary
803/// derivative stacks contribute only their value entry, while structured
804/// quadratic operators retain their representation-specific primal lowering.
805#[derive(Clone, Copy, Debug, PartialEq)]
806pub struct RuntimeValue {
807    value: f64,
808    dimension: usize,
809}
810
811impl<'arena> RuntimeJetScalar<'arena> for RuntimeValue {
812    type Workspace = ();
813
814    #[inline(always)]
815    fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
816        Self {
817            value: c,
818            dimension,
819        }
820    }
821
822    #[inline(always)]
823    fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
824        assert!(
825            axis < dimension,
826            "runtime value variable axis out of bounds"
827        );
828        Self {
829            value: x,
830            dimension,
831        }
832    }
833
834    #[inline(always)]
835    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
836        inputs: &[Self],
837        coefficients: &C,
838        dimension: usize,
839        &(): &'arena Self::Workspace,
840    ) -> Self {
841        assert_eq!(inputs.len(), coefficients.dimension());
842        assert!(inputs.iter().all(|input| input.dimension == dimension));
843        Self {
844            value: coefficients.quadratic_value(inputs, |input| input.value),
845            dimension,
846        }
847    }
848
849    #[inline(always)]
850    fn linear_combination(
851        inputs: &[Self],
852        weights: &[f64],
853        dimension: usize,
854        &(): &'arena Self::Workspace,
855    ) -> Self {
856        assert_eq!(inputs.len(), weights.len());
857        assert!(inputs.iter().all(|input| input.dimension == dimension));
858        let value = inputs
859            .iter()
860            .zip(weights)
861            .map(|(input, &weight)| input.value * weight)
862            .sum();
863        Self { value, dimension }
864    }
865
866    #[inline(always)]
867    fn add_constant(&self, constant: f64, &(): &'arena Self::Workspace) -> Self {
868        Self {
869            value: self.value + constant,
870            dimension: self.dimension,
871        }
872    }
873
874    #[inline(always)]
875    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
876        self.assert_same_dimension(right);
877        self.assert_same_dimension(addend);
878        Self {
879            value: self.value * right.value + addend.value,
880            dimension: self.dimension,
881        }
882    }
883
884    #[inline(always)]
885    fn composed_sum(
886        inputs: &[Self],
887        derivative_stacks: &[[f64; 5]],
888        dimension: usize,
889        &(): &'arena Self::Workspace,
890    ) -> Self {
891        assert_eq!(inputs.len(), derivative_stacks.len());
892        assert!(inputs.iter().all(|input| input.dimension == dimension));
893        Self {
894            value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
895            dimension,
896        }
897    }
898
899    #[inline(always)]
900    fn product(&self, right: &Self) -> Self {
901        self.mul(right)
902    }
903
904    #[inline(always)]
905    fn affine_compose(
906        &self,
907        _: f64,
908        _: f64,
909        derivative_stack: [f64; 5],
910        &(): &'arena Self::Workspace,
911    ) -> Self {
912        Self {
913            value: derivative_stack[0],
914            dimension: self.dimension,
915        }
916    }
917
918    #[inline(always)]
919    fn affine_composed_sum(
920        inputs: &[Self],
921        input_scales: &[f64],
922        derivative_stacks: &[[f64; 5]],
923        dimension: usize,
924        &(): &'arena Self::Workspace,
925    ) -> Self {
926        assert_eq!(inputs.len(), input_scales.len());
927        assert_eq!(inputs.len(), derivative_stacks.len());
928        assert!(inputs.iter().all(|input| input.dimension == dimension));
929        Self {
930            value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
931            dimension,
932        }
933    }
934
935    #[inline(always)]
936    fn dimension(&self) -> usize {
937        self.dimension
938    }
939
940    #[inline(always)]
941    fn value(&self) -> f64 {
942        self.value
943    }
944
945    #[inline(always)]
946    fn add(&self, other: &Self) -> Self {
947        self.assert_same_dimension(other);
948        Self {
949            value: self.value + other.value,
950            dimension: self.dimension,
951        }
952    }
953
954    #[inline(always)]
955    fn sub(&self, other: &Self) -> Self {
956        self.assert_same_dimension(other);
957        Self {
958            value: self.value - other.value,
959            dimension: self.dimension,
960        }
961    }
962
963    #[inline(always)]
964    fn mul(&self, other: &Self) -> Self {
965        self.assert_same_dimension(other);
966        Self {
967            value: self.value * other.value,
968            dimension: self.dimension,
969        }
970    }
971
972    #[inline(always)]
973    fn neg(&self) -> Self {
974        Self {
975            value: -self.value,
976            dimension: self.dimension,
977        }
978    }
979
980    #[inline(always)]
981    fn scale(&self, scale: f64) -> Self {
982        Self {
983            value: self.value * scale,
984            dimension: self.dimension,
985        }
986    }
987
988    #[inline(always)]
989    fn compose_unary(&self, derivative_stack: [f64; 5]) -> Self {
990        Self {
991            value: derivative_stack[0],
992            dimension: self.dimension,
993        }
994    }
995}
996
997impl RuntimeValue {
998    #[inline(always)]
999    fn assert_same_dimension(&self, other: &Self) {
1000        assert_eq!(self.dimension, other.dimension);
1001    }
1002}
1003
1004/// Adapter that presents any const-generic [`JetScalar<K>`] through the
1005/// runtime-dimension interface.  It is used by derivative oracles so the same
1006/// row program can be instantiated at a fixed tower and at a dynamic packed
1007/// scalar; production code unwraps the inner fixed tower after evaluation.
1008#[derive(Clone, Copy, Debug)]
1009#[repr(transparent)]
1010pub struct FixedRuntimeJet<S, const K: usize> {
1011    inner: S,
1012}
1013
1014impl<S, const K: usize> FixedRuntimeJet<S, K> {
1015    /// Wrap a fixed-dimension scalar so a runtime-dimension row expression can
1016    /// evaluate it without authoring a second const-generic body.
1017    #[inline(always)]
1018    #[must_use]
1019    pub fn from_inner(inner: S) -> Self {
1020        Self { inner }
1021    }
1022
1023    /// Recover the wrapped const-generic scalar.
1024    #[inline(always)]
1025    #[must_use]
1026    pub fn into_inner(self) -> S {
1027        self.inner
1028    }
1029}
1030
1031impl<'arena, S: JetScalar<K>, const K: usize> RuntimeJetScalar<'arena> for FixedRuntimeJet<S, K> {
1032    type Workspace = ();
1033
1034    #[inline(always)]
1035    fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
1036        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1037        Self {
1038            inner: S::constant(c),
1039        }
1040    }
1041
1042    #[inline(always)]
1043    fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
1044        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1045        Self {
1046            inner: S::variable(x, axis),
1047        }
1048    }
1049
1050    #[inline(always)]
1051    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
1052        inputs: &[Self],
1053        coefficients: &C,
1054        dimension: usize,
1055        &(): &'arena Self::Workspace,
1056    ) -> Self {
1057        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1058        assert_eq!(inputs.len(), coefficients.dimension());
1059        // SAFETY: `FixedRuntimeJet<S, K>` is `repr(transparent)` with exactly one
1060        // non-zero-sized field of type `S`. Therefore each element has the same
1061        // layout and alignment as `S`; the cast preserves the allocation,
1062        // provenance, element count, and shared lifetime of `inputs`, and no
1063        // mutable reference is created while the original slice is borrowed.
1064        let inner =
1065            unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1066        Self {
1067            inner: S::symmetric_quadratic_form(inner, coefficients),
1068        }
1069    }
1070
1071    #[inline(always)]
1072    fn linear_combination(
1073        inputs: &[Self],
1074        weights: &[f64],
1075        dimension: usize,
1076        &(): &'arena Self::Workspace,
1077    ) -> Self {
1078        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1079        assert_eq!(inputs.len(), weights.len());
1080        // SAFETY: `FixedRuntimeJet<S, K>` is `repr(transparent)` over its sole
1081        // non-zero-sized `S` field, so the shared slice has identical element
1082        // layout, alignment, provenance, length, and lifetime after the cast.
1083        let inner =
1084            unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1085        Self {
1086            inner: S::linear_combination(inner, weights),
1087        }
1088    }
1089
1090    #[inline(always)]
1091    fn add_constant(&self, constant: f64, &(): &'arena Self::Workspace) -> Self {
1092        Self {
1093            inner: self.inner.add_constant(constant),
1094        }
1095    }
1096
1097    #[inline(always)]
1098    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
1099        Self {
1100            inner: self.inner.multiply_add(&right.inner, &addend.inner),
1101        }
1102    }
1103
1104    #[inline(always)]
1105    fn composed_sum(
1106        inputs: &[Self],
1107        derivative_stacks: &[[f64; 5]],
1108        dimension: usize,
1109        &(): &'arena Self::Workspace,
1110    ) -> Self {
1111        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1112        // SAFETY: `FixedRuntimeJet<S, K>` is `repr(transparent)` over its sole
1113        // non-zero-sized `S` field, so the shared slice has identical element
1114        // layout, alignment, provenance, length, and lifetime after the cast.
1115        let inner =
1116            unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1117        Self {
1118            inner: S::composed_sum(inner, derivative_stacks),
1119        }
1120    }
1121
1122    #[inline(always)]
1123    fn product(&self, right: &Self) -> Self {
1124        Self {
1125            inner: self.inner.product(&right.inner),
1126        }
1127    }
1128
1129    #[inline(always)]
1130    fn affine_compose(
1131        &self,
1132        input_scale: f64,
1133        input_shift: f64,
1134        derivative_stack: [f64; 5],
1135        &(): &'arena Self::Workspace,
1136    ) -> Self {
1137        Self {
1138            inner: self
1139                .inner
1140                .affine_compose(input_scale, input_shift, derivative_stack),
1141        }
1142    }
1143
1144    #[inline(always)]
1145    fn affine_composed_sum(
1146        inputs: &[Self],
1147        input_scales: &[f64],
1148        derivative_stacks: &[[f64; 5]],
1149        dimension: usize,
1150        &(): &'arena Self::Workspace,
1151    ) -> Self {
1152        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1153        // SAFETY: `FixedRuntimeJet<S, K>` is `repr(transparent)` over its sole
1154        // non-zero-sized `S` field, so the shared slice has identical element
1155        // layout, alignment, provenance, length, and lifetime after the cast.
1156        let inner =
1157            unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1158        Self {
1159            inner: S::affine_composed_sum(inner, input_scales, derivative_stacks),
1160        }
1161    }
1162
1163    #[inline(always)]
1164    fn shared_multiply_add_affine_composed_sum<const N: usize>(
1165        lefts: &[&Self; N],
1166        right: &Self,
1167        addend: &Self,
1168        addend_scales: &[f64; N],
1169        input_scales: &[f64; N],
1170        derivative_stacks: &[[f64; 5]; N],
1171        dimension: usize,
1172        &(): &'arena Self::Workspace,
1173    ) -> Self {
1174        assert_eq!(dimension, K, "fixed jet dimension mismatch");
1175        let left_inner: [&S; N] = std::array::from_fn(|term| &lefts[term].inner);
1176        Self {
1177            inner: S::shared_multiply_add_affine_composed_sum(
1178                &left_inner,
1179                &right.inner,
1180                &addend.inner,
1181                addend_scales,
1182                input_scales,
1183                derivative_stacks,
1184            ),
1185        }
1186    }
1187
1188    #[inline(always)]
1189    fn dimension(&self) -> usize {
1190        K
1191    }
1192
1193    #[inline(always)]
1194    fn value(&self) -> f64 {
1195        self.inner.value()
1196    }
1197
1198    #[inline(always)]
1199    fn add(&self, o: &Self) -> Self {
1200        Self {
1201            inner: self.inner.add(&o.inner),
1202        }
1203    }
1204
1205    #[inline(always)]
1206    fn sub(&self, o: &Self) -> Self {
1207        Self {
1208            inner: self.inner.sub(&o.inner),
1209        }
1210    }
1211
1212    #[inline(always)]
1213    fn mul(&self, o: &Self) -> Self {
1214        Self {
1215            inner: self.inner.mul(&o.inner),
1216        }
1217    }
1218
1219    #[inline(always)]
1220    fn neg(&self) -> Self {
1221        Self {
1222            inner: self.inner.neg(),
1223        }
1224    }
1225
1226    #[inline(always)]
1227    fn scale(&self, s: f64) -> Self {
1228        Self {
1229            inner: self.inner.scale(s),
1230        }
1231    }
1232
1233    #[inline(always)]
1234    fn compose_unary(&self, d: [f64; 5]) -> Self {
1235        Self {
1236            inner: self.inner.compose_unary(d),
1237        }
1238    }
1239}
1240
1241/// Reusable storage for runtime-sized packed jets. Scalar primitives write
1242/// into this bump arena, so arithmetic performs no heap allocation. Callers
1243/// reserve once per worker/chunk and [`reset`](Self::reset) between rows.
1244#[derive(Debug)]
1245pub struct DynamicJetArena {
1246    bump: bumpalo::Bump,
1247}
1248
1249impl DynamicJetArena {
1250    /// Create an arena with the allocator's default initial chunk.
1251    #[must_use]
1252    pub fn new() -> Self {
1253        Self {
1254            bump: bumpalo::Bump::new(),
1255        }
1256    }
1257
1258    /// Create an arena with a row-program-selected initial byte capacity.
1259    #[must_use]
1260    pub fn with_capacity(bytes: usize) -> Self {
1261        Self {
1262            bump: bumpalo::Bump::with_capacity(bytes),
1263        }
1264    }
1265
1266    /// Reclaim all scalar outputs and compact a fragmented high-water mark
1267    /// into one retained chunk.
1268    ///
1269    /// `bumpalo::Bump::reset` keeps only its newest chunk and returns every
1270    /// older chunk to the global allocator. Runtime jet tapes routinely span
1271    /// several geometrically-grown chunks, so a bare reset would allocate
1272    /// those discarded chunks again on every row. When reset exposes that
1273    /// fragmentation, replace the empty arena once with a single chunk large
1274    /// enough for the complete prior tape. Later equal-or-smaller rows reuse
1275    /// that chunk without allocator traffic.
1276    pub fn reset(&mut self) {
1277        let high_water = self.bump.allocated_bytes();
1278        self.bump.reset();
1279        if self.bump.allocated_bytes() < high_water {
1280            self.bump = bumpalo::Bump::with_capacity(high_water);
1281        }
1282    }
1283
1284    /// Bytes currently reserved from the global allocator. A warm-reset-warm
1285    /// benchmark uses this to prove the second row requires no arena growth.
1286    #[must_use]
1287    pub fn allocated_bytes(&self) -> usize {
1288        self.bump.allocated_bytes()
1289    }
1290
1291    #[inline(always)]
1292    fn zeros(&self, len: usize) -> &mut [f64] {
1293        self.bump.alloc_slice_fill_copy(len, 0.0)
1294    }
1295
1296    /// Allocate and initialize a runtime-sized slice in the arena. Row programs
1297    /// use this for their primary-scalar arrays so those arrays share the same
1298    /// reusable workspace as derivative channels.
1299    #[inline(always)]
1300    pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
1301        self.bump.alloc_slice_fill_with(len, fill)
1302    }
1303}
1304
1305impl Default for DynamicJetArena {
1306    fn default() -> Self {
1307        Self::new()
1308    }
1309}
1310
1311/// Runtime-sized packed first-order scalar: value plus arena-backed gradient.
1312#[derive(Clone, Copy, Debug)]
1313pub struct DynamicOrder1<'arena> {
1314    arena: &'arena DynamicJetArena,
1315    /// Value channel.
1316    pub v: f64,
1317    /// Gradient channel, length [`Self::dimension`].
1318    pub g: &'arena [f64],
1319}
1320
1321impl DynamicOrder1<'_> {
1322    /// Gradient channel.
1323    #[inline]
1324    #[must_use]
1325    pub fn g(&self) -> &[f64] {
1326        self.g
1327    }
1328
1329    #[inline]
1330    fn assert_compatible(&self, o: &Self) {
1331        assert_eq!(
1332            self.g.len(),
1333            o.g.len(),
1334            "dynamic first-order jet dimension mismatch"
1335        );
1336        assert!(
1337            std::ptr::eq(self.arena, o.arena),
1338            "dynamic jets belong to different arenas"
1339        );
1340    }
1341}
1342
1343impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder1<'arena> {
1344    type Workspace = DynamicJetArena;
1345
1346    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1347        Self {
1348            arena,
1349            v: c,
1350            g: arena.zeros(dimension),
1351        }
1352    }
1353
1354    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1355        assert!(
1356            axis < dimension,
1357            "dynamic first-order jet axis out of bounds"
1358        );
1359        let g = arena.zeros(dimension);
1360        g[axis] = 1.0;
1361        Self { arena, v: x, g }
1362    }
1363
1364    fn dimension(&self) -> usize {
1365        self.g.len()
1366    }
1367    fn value(&self) -> f64 {
1368        self.v
1369    }
1370
1371    fn add(&self, o: &Self) -> Self {
1372        self.assert_compatible(o);
1373        let g = self.arena.zeros(self.dimension());
1374        for i in 0..g.len() {
1375            g[i] = self.g[i] + o.g[i];
1376        }
1377        Self {
1378            arena: self.arena,
1379            v: self.v + o.v,
1380            g,
1381        }
1382    }
1383
1384    fn sub(&self, o: &Self) -> Self {
1385        self.assert_compatible(o);
1386        let g = self.arena.zeros(self.dimension());
1387        for i in 0..g.len() {
1388            g[i] = self.g[i] - o.g[i];
1389        }
1390        Self {
1391            arena: self.arena,
1392            v: self.v - o.v,
1393            g,
1394        }
1395    }
1396
1397    fn mul(&self, o: &Self) -> Self {
1398        self.assert_compatible(o);
1399        let g = self.arena.zeros(self.dimension());
1400        for i in 0..g.len() {
1401            g[i] = self.v * o.g[i] + self.g[i] * o.v;
1402        }
1403        Self {
1404            arena: self.arena,
1405            v: self.v * o.v,
1406            g,
1407        }
1408    }
1409
1410    fn neg(&self) -> Self {
1411        self.scale(-1.0)
1412    }
1413
1414    fn scale(&self, s: f64) -> Self {
1415        let g = self.arena.zeros(self.dimension());
1416        for i in 0..g.len() {
1417            g[i] = self.g[i] * s;
1418        }
1419        Self {
1420            arena: self.arena,
1421            v: self.v * s,
1422            g,
1423        }
1424    }
1425
1426    fn compose_unary(&self, d: [f64; 5]) -> Self {
1427        let g = self.arena.zeros(self.dimension());
1428        for i in 0..g.len() {
1429            g[i] = d[1] * self.g[i];
1430        }
1431        Self {
1432            arena: self.arena,
1433            v: d[0],
1434            g,
1435        }
1436    }
1437}
1438
1439/// Runtime-sized packed second-order scalar: value, gradient, and a row-major
1440/// Hessian. Storage is `O(K^2)` in the row's actual primary dimension and comes
1441/// from the row's reusable [`DynamicJetArena`].
1442#[derive(Clone, Copy, Debug)]
1443pub struct DynamicOrder2<'arena> {
1444    arena: &'arena DynamicJetArena,
1445    /// Value channel.
1446    pub v: f64,
1447    /// Gradient channel.
1448    pub g: &'arena [f64],
1449    /// Row-major Hessian channel.
1450    pub h: &'arena [f64],
1451}
1452
1453impl DynamicOrder2<'_> {
1454    /// Construct an arena-backed second-order scalar from channel functions.
1455    ///
1456    /// This is the allocation-free extension seam for exact projected products
1457    /// whose channel law is not ordinary scalar multiplication. Each channel is
1458    /// written directly into the row arena, so downstream runtime-jet programs
1459    /// can keep their specialized algebra without materializing temporary
1460    /// `Vec`s or exposing the arena pointer stored by [`DynamicOrder2`]. The
1461    /// Hessian function is evaluated on the upper triangle and mirrored because
1462    /// a scalar Hessian is symmetric.
1463    #[inline]
1464    #[must_use]
1465    pub fn from_channel_functions<'arena>(
1466        value: f64,
1467        dimension: usize,
1468        arena: &'arena DynamicJetArena,
1469        mut gradient: impl FnMut(usize) -> f64,
1470        mut hessian: impl FnMut(usize, usize) -> f64,
1471    ) -> DynamicOrder2<'arena> {
1472        let g = arena.alloc_slice_fill_with(dimension, |axis| gradient(axis));
1473        let h = arena.zeros(dimension * dimension);
1474        for row in 0..dimension {
1475            for column in row..dimension {
1476                let channel = hessian(row, column);
1477                h[row * dimension + column] = channel;
1478                h[column * dimension + row] = channel;
1479            }
1480        }
1481        DynamicOrder2 {
1482            arena,
1483            v: value,
1484            g,
1485            h,
1486        }
1487    }
1488
1489    /// Gradient channel.
1490    #[inline]
1491    #[must_use]
1492    pub fn g(&self) -> &[f64] {
1493        self.g
1494    }
1495
1496    /// Row-major Hessian channel.
1497    #[inline]
1498    #[must_use]
1499    pub fn h(&self) -> &[f64] {
1500        self.h
1501    }
1502
1503    /// Hessian entry `(row, col)`.
1504    #[inline]
1505    #[must_use]
1506    pub fn h_at(&self, row: usize, col: usize) -> f64 {
1507        self.h[row * self.dimension() + col]
1508    }
1509
1510    #[inline(always)]
1511    fn assert_compatible(&self, o: &Self) {
1512        assert_eq!(
1513            self.g.len(),
1514            o.g.len(),
1515            "dynamic second-order jet dimension mismatch"
1516        );
1517        assert_eq!(
1518            self.h.len(),
1519            o.h.len(),
1520            "dynamic second-order jet Hessian mismatch"
1521        );
1522        assert!(
1523            std::ptr::eq(self.arena, o.arena),
1524            "dynamic jets belong to different arenas"
1525        );
1526    }
1527}
1528
1529impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder2<'arena> {
1530    type Workspace = DynamicJetArena;
1531
1532    #[inline(always)]
1533    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1534        Self {
1535            arena,
1536            v: c,
1537            g: arena.zeros(dimension),
1538            h: arena.zeros(dimension * dimension),
1539        }
1540    }
1541
1542    #[inline(always)]
1543    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1544        assert!(
1545            axis < dimension,
1546            "dynamic second-order jet axis out of bounds"
1547        );
1548        let g = arena.zeros(dimension);
1549        g[axis] = 1.0;
1550        Self {
1551            arena,
1552            v: x,
1553            g,
1554            h: arena.zeros(dimension * dimension),
1555        }
1556    }
1557
1558    #[inline(always)]
1559    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
1560        inputs: &[Self],
1561        coefficients: &C,
1562        dimension: usize,
1563        arena: &'arena DynamicJetArena,
1564    ) -> Self {
1565        assert_eq!(inputs.len(), coefficients.dimension());
1566        assert!(
1567            inputs.iter().all(|input| {
1568                input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1569            }),
1570            "dynamic quadratic-form jets must share dimension and arena"
1571        );
1572        let input_dimension = inputs.len();
1573        let values = arena.zeros(input_dimension);
1574        for (value, input) in values.iter_mut().zip(inputs) {
1575            *value = input.v;
1576        }
1577        let projected = arena.zeros(input_dimension);
1578        coefficients.multiply(values, projected);
1579
1580        let mut value = 0.0;
1581        for axis in 0..input_dimension {
1582            value += values[axis] * projected[axis];
1583        }
1584        let gradient = arena.zeros(dimension);
1585        for primary in 0..dimension {
1586            let mut channel = 0.0;
1587            for axis in 0..input_dimension {
1588                channel += projected[axis] * inputs[axis].g[primary];
1589            }
1590            gradient[primary] = 2.0 * channel;
1591        }
1592        let hessian = arena.zeros(dimension * dimension);
1593        let input_gradient = arena.zeros(input_dimension);
1594        let projected_gradient = arena.zeros(input_dimension);
1595        for primary_b in 0..dimension {
1596            for row in 0..input_dimension {
1597                input_gradient[row] = inputs[row].g[primary_b];
1598            }
1599            coefficients.multiply(input_gradient, projected_gradient);
1600            for primary_a in 0..=primary_b {
1601                let mut inherited = 0.0;
1602                let mut curvature = 0.0;
1603                for row in 0..input_dimension {
1604                    inherited += projected[row] * inputs[row].h[primary_a * dimension + primary_b];
1605                    curvature += inputs[row].g[primary_a] * projected_gradient[row];
1606                }
1607                let channel = 2.0 * (inherited + curvature);
1608                hessian[primary_a * dimension + primary_b] = channel;
1609                hessian[primary_b * dimension + primary_a] = channel;
1610            }
1611        }
1612        Self {
1613            arena,
1614            v: value,
1615            g: gradient,
1616            h: hessian,
1617        }
1618    }
1619
1620    #[inline(always)]
1621    fn product(&self, right: &Self) -> Self {
1622        self.assert_compatible(right);
1623        let dimension = self.dimension();
1624        let gradient = self.arena.zeros(dimension);
1625        let hessian = self.arena.zeros(dimension * dimension);
1626        for primary in 0..dimension {
1627            gradient[primary] = self.v * right.g[primary] + self.g[primary] * right.v;
1628            for other in primary..dimension {
1629                let index = primary * dimension + other;
1630                let channel = self.v * right.h[index]
1631                    + self.g[primary] * right.g[other]
1632                    + self.g[other] * right.g[primary]
1633                    + self.h[index] * right.v;
1634                hessian[index] = channel;
1635                hessian[other * dimension + primary] = channel;
1636            }
1637        }
1638        Self {
1639            arena: self.arena,
1640            v: self.v * right.v,
1641            g: gradient,
1642            h: hessian,
1643        }
1644    }
1645
1646    #[inline(always)]
1647    fn affine_compose(
1648        &self,
1649        input_scale: f64,
1650        input_shift: f64,
1651        derivative_stack: [f64; 5],
1652        arena: &'arena DynamicJetArena,
1653    ) -> Self {
1654        assert!(std::ptr::eq(self.arena, arena));
1655        assert!(input_shift.is_finite(), "affine input shift must be finite");
1656        let dimension = self.dimension();
1657        let first = derivative_stack[1] * input_scale;
1658        let second = derivative_stack[2] * input_scale * input_scale;
1659        let gradient = arena.zeros(dimension);
1660        let hessian = arena.zeros(dimension * dimension);
1661        for primary in 0..dimension {
1662            gradient[primary] = first * self.g[primary];
1663            for other in primary..dimension {
1664                let index = primary * dimension + other;
1665                let channel = first * self.h[index] + second * self.g[primary] * self.g[other];
1666                hessian[index] = channel;
1667                hessian[other * dimension + primary] = channel;
1668            }
1669        }
1670        Self {
1671            arena,
1672            v: derivative_stack[0],
1673            g: gradient,
1674            h: hessian,
1675        }
1676    }
1677
1678    #[inline(always)]
1679    fn affine_composed_sum(
1680        inputs: &[Self],
1681        input_scales: &[f64],
1682        derivative_stacks: &[[f64; 5]],
1683        dimension: usize,
1684        arena: &'arena DynamicJetArena,
1685    ) -> Self {
1686        assert_eq!(inputs.len(), input_scales.len());
1687        assert_eq!(inputs.len(), derivative_stacks.len());
1688        assert!(
1689            inputs.iter().all(|input| {
1690                input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1691            }),
1692            "dynamic affine-composed-sum jets must share dimension and arena"
1693        );
1694        let gradient = arena.zeros(dimension);
1695        let hessian = arena.zeros(dimension * dimension);
1696        let mut value = 0.0;
1697        for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
1698        {
1699            let first = stack[1] * input_scale;
1700            let second = stack[2] * input_scale * input_scale;
1701            value += stack[0];
1702            for primary in 0..dimension {
1703                gradient[primary] += first * input.g[primary];
1704                for other in primary..dimension {
1705                    let index = primary * dimension + other;
1706                    hessian[index] +=
1707                        first * input.h[index] + second * input.g[primary] * input.g[other];
1708                }
1709            }
1710        }
1711        for primary in 0..dimension {
1712            for other in primary + 1..dimension {
1713                hessian[other * dimension + primary] = hessian[primary * dimension + other];
1714            }
1715        }
1716        Self {
1717            arena,
1718            v: value,
1719            g: gradient,
1720            h: hessian,
1721        }
1722    }
1723
1724    #[inline(always)]
1725    fn shared_multiply_add_affine_composed_sum<const N: usize>(
1726        lefts: &[&Self; N],
1727        right: &Self,
1728        addend: &Self,
1729        addend_scales: &[f64; N],
1730        input_scales: &[f64; N],
1731        derivative_stacks: &[[f64; 5]; N],
1732        dimension: usize,
1733        arena: &'arena DynamicJetArena,
1734    ) -> Self {
1735        assert!(
1736            lefts.iter().all(|input| {
1737                input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1738            }) && (N == 0 || (right.dimension() == dimension && std::ptr::eq(right.arena, arena))),
1739            "dynamic fused product-composition jets must share dimension and arena"
1740        );
1741        let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
1742        assert!(
1743            !addend_live || (addend.dimension() == dimension && std::ptr::eq(addend.arena, arena)),
1744            "live dynamic fused addends must share dimension and arena"
1745        );
1746        let (representatives, term_sources, source_count) =
1747            canonical_shared_source_schedule::<N>(|term, representative| {
1748                std::ptr::eq(lefts[term], lefts[representative])
1749                    && addend_scales[term] == addend_scales[representative]
1750            });
1751        let (value, source_derivatives) =
1752            aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
1753        let source_gradients = arena.zeros(source_count * dimension);
1754        let gradient = arena.zeros(dimension);
1755        let hessian = arena.zeros(dimension * dimension);
1756        let mut right_first = 0.0;
1757        let mut addend_first = 0.0;
1758        for source in 0..source_count {
1759            let term = representatives[source];
1760            let first = source_derivatives[source][1];
1761            right_first += first * lefts[term].v;
1762            addend_first += first * addend_scales[term];
1763            for primary in 0..dimension {
1764                let product_gradient =
1765                    lefts[term].v * right.g[primary] + lefts[term].g[primary] * right.v;
1766                let inner_gradient = if addend_scales[term] == 0.0 {
1767                    product_gradient
1768                } else if addend_scales[term] == 1.0 {
1769                    product_gradient + addend.g[primary]
1770                } else {
1771                    product_gradient + addend_scales[term] * addend.g[primary]
1772                };
1773                source_gradients[source * dimension + primary] = inner_gradient;
1774                gradient[primary] += first * lefts[term].g[primary] * right.v;
1775            }
1776        }
1777        if N != 0 {
1778            for primary in 0..dimension {
1779                gradient[primary] += right_first * right.g[primary];
1780            }
1781        }
1782        if addend_live {
1783            for primary in 0..dimension {
1784                gradient[primary] += addend_first * addend.g[primary];
1785            }
1786        }
1787        for primary in 0..dimension {
1788            for other in primary..dimension {
1789                let index = primary * dimension + other;
1790                let mut channel = if N == 0 {
1791                    0.0
1792                } else {
1793                    right_first * right.h[index]
1794                };
1795                if addend_live {
1796                    channel += addend_first * addend.h[index];
1797                }
1798                for source in 0..source_count {
1799                    let term = representatives[source];
1800                    let local_product_hessian = lefts[term].g[primary] * right.g[other]
1801                        + lefts[term].g[other] * right.g[primary]
1802                        + lefts[term].h[index] * right.v;
1803                    channel += source_derivatives[source][1] * local_product_hessian
1804                        + source_derivatives[source][2]
1805                            * source_gradients[source * dimension + primary]
1806                            * source_gradients[source * dimension + other];
1807                }
1808                hessian[index] = channel;
1809                hessian[other * dimension + primary] = channel;
1810            }
1811        }
1812        Self {
1813            arena,
1814            v: value,
1815            g: gradient,
1816            h: hessian,
1817        }
1818    }
1819
1820    #[inline(always)]
1821    fn add_constant(&self, constant: f64, arena: &'arena DynamicJetArena) -> Self {
1822        assert!(std::ptr::eq(self.arena, arena));
1823        Self {
1824            arena,
1825            v: self.v + constant,
1826            g: self.g,
1827            h: self.h,
1828        }
1829    }
1830
1831    #[inline(always)]
1832    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
1833        self.assert_compatible(right);
1834        self.assert_compatible(addend);
1835        let dimension = self.dimension();
1836        let gradient = self.arena.zeros(dimension);
1837        let hessian = self.arena.zeros(dimension * dimension);
1838        for primary in 0..dimension {
1839            gradient[primary] =
1840                self.v * right.g[primary] + self.g[primary] * right.v + addend.g[primary];
1841            for other in primary..dimension {
1842                let index = primary * dimension + other;
1843                let channel = self.v * right.h[index]
1844                    + self.g[primary] * right.g[other]
1845                    + self.g[other] * right.g[primary]
1846                    + self.h[index] * right.v
1847                    + addend.h[index];
1848                hessian[index] = channel;
1849                hessian[other * dimension + primary] = channel;
1850            }
1851        }
1852        Self {
1853            arena: self.arena,
1854            v: self.v * right.v + addend.v,
1855            g: gradient,
1856            h: hessian,
1857        }
1858    }
1859
1860    #[inline(always)]
1861    fn composed_sum(
1862        inputs: &[Self],
1863        derivative_stacks: &[[f64; 5]],
1864        dimension: usize,
1865        arena: &'arena DynamicJetArena,
1866    ) -> Self {
1867        assert_eq!(inputs.len(), derivative_stacks.len());
1868        assert!(
1869            inputs.iter().all(|input| {
1870                input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1871            }),
1872            "dynamic composed-sum jets must share dimension and arena"
1873        );
1874        let gradient = arena.zeros(dimension);
1875        let hessian = arena.zeros(dimension * dimension);
1876        let mut value = 0.0;
1877        for (input, stack) in inputs.iter().zip(derivative_stacks) {
1878            value += stack[0];
1879            for primary in 0..dimension {
1880                gradient[primary] += stack[1] * input.g[primary];
1881                for other in primary..dimension {
1882                    let index = primary * dimension + other;
1883                    hessian[index] +=
1884                        stack[1] * input.h[index] + stack[2] * input.g[primary] * input.g[other];
1885                }
1886            }
1887        }
1888        for primary in 0..dimension {
1889            for other in primary + 1..dimension {
1890                hessian[other * dimension + primary] = hessian[primary * dimension + other];
1891            }
1892        }
1893        Self {
1894            arena,
1895            v: value,
1896            g: gradient,
1897            h: hessian,
1898        }
1899    }
1900
1901    #[inline(always)]
1902    fn linear_combination(
1903        inputs: &[Self],
1904        weights: &[f64],
1905        dimension: usize,
1906        arena: &'arena DynamicJetArena,
1907    ) -> Self {
1908        assert_eq!(inputs.len(), weights.len());
1909        assert!(
1910            inputs.iter().all(|input| {
1911                input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1912            }),
1913            "dynamic linear-combination jets must share dimension and arena"
1914        );
1915        let mut value = 0.0;
1916        for (input, &weight) in inputs.iter().zip(weights) {
1917            value += input.v * weight;
1918        }
1919        let gradient = arena.zeros(dimension);
1920        let hessian = arena.zeros(dimension * dimension);
1921        for primary in 0..dimension {
1922            for (input, &weight) in inputs.iter().zip(weights) {
1923                gradient[primary] += input.g[primary] * weight;
1924            }
1925            for other in primary..dimension {
1926                let index = primary * dimension + other;
1927                for (input, &weight) in inputs.iter().zip(weights) {
1928                    hessian[index] += input.h[index] * weight;
1929                }
1930                hessian[other * dimension + primary] = hessian[index];
1931            }
1932        }
1933        Self {
1934            arena,
1935            v: value,
1936            g: gradient,
1937            h: hessian,
1938        }
1939    }
1940
1941    #[inline(always)]
1942    fn dimension(&self) -> usize {
1943        self.g.len()
1944    }
1945
1946    #[inline(always)]
1947    fn value(&self) -> f64 {
1948        self.v
1949    }
1950
1951    #[inline(always)]
1952    fn add(&self, o: &Self) -> Self {
1953        self.assert_compatible(o);
1954        let dimension = self.dimension();
1955        let g = self.arena.zeros(dimension);
1956        let h = self.arena.zeros(self.h.len());
1957        for i in 0..g.len() {
1958            g[i] = self.g[i] + o.g[i];
1959        }
1960        for row in 0..dimension {
1961            for column in row..dimension {
1962                let index = row * dimension + column;
1963                let channel = self.h[index] + o.h[index];
1964                h[index] = channel;
1965                h[column * dimension + row] = channel;
1966            }
1967        }
1968        Self {
1969            arena: self.arena,
1970            v: self.v + o.v,
1971            g,
1972            h,
1973        }
1974    }
1975
1976    #[inline(always)]
1977    fn sub(&self, o: &Self) -> Self {
1978        self.assert_compatible(o);
1979        let dimension = self.dimension();
1980        let g = self.arena.zeros(dimension);
1981        let h = self.arena.zeros(self.h.len());
1982        for i in 0..g.len() {
1983            g[i] = self.g[i] - o.g[i];
1984        }
1985        for row in 0..dimension {
1986            for column in row..dimension {
1987                let index = row * dimension + column;
1988                let channel = self.h[index] - o.h[index];
1989                h[index] = channel;
1990                h[column * dimension + row] = channel;
1991            }
1992        }
1993        Self {
1994            arena: self.arena,
1995            v: self.v - o.v,
1996            g,
1997            h,
1998        }
1999    }
2000
2001    #[inline(always)]
2002    fn mul(&self, o: &Self) -> Self {
2003        self.assert_compatible(o);
2004        let n = self.dimension();
2005        let g = self.arena.zeros(n);
2006        let h = self.arena.zeros(n * n);
2007        for i in 0..n {
2008            g[i] = self.v * o.g[i] + self.g[i] * o.v;
2009        }
2010        for i in 0..n {
2011            for j in i..n {
2012                let ij = i * n + j;
2013                let hij =
2014                    self.v * o.h[ij] + self.g[i] * o.g[j] + self.g[j] * o.g[i] + self.h[ij] * o.v;
2015                h[ij] = hij;
2016                h[j * n + i] = hij;
2017            }
2018        }
2019        Self {
2020            arena: self.arena,
2021            v: self.v * o.v,
2022            g,
2023            h,
2024        }
2025    }
2026
2027    #[inline(always)]
2028    fn neg(&self) -> Self {
2029        self.scale(-1.0)
2030    }
2031
2032    #[inline(always)]
2033    fn scale(&self, s: f64) -> Self {
2034        let dimension = self.dimension();
2035        let g = self.arena.zeros(dimension);
2036        let h = self.arena.zeros(self.h.len());
2037        for i in 0..g.len() {
2038            g[i] = self.g[i] * s;
2039        }
2040        for row in 0..dimension {
2041            for column in row..dimension {
2042                let index = row * dimension + column;
2043                let channel = self.h[index] * s;
2044                h[index] = channel;
2045                h[column * dimension + row] = channel;
2046            }
2047        }
2048        Self {
2049            arena: self.arena,
2050            v: self.v * s,
2051            g,
2052            h,
2053        }
2054    }
2055
2056    #[inline(always)]
2057    fn compose_unary(&self, d: [f64; 5]) -> Self {
2058        let n = self.dimension();
2059        let g = self.arena.zeros(n);
2060        let h = self.arena.zeros(n * n);
2061        for i in 0..n {
2062            g[i] = d[1] * self.g[i];
2063        }
2064        for i in 0..n {
2065            for j in i..n {
2066                let ij = i * n + j;
2067                let channel = d[1] * self.h[ij] + d[2] * self.g[i] * self.g[j];
2068                h[ij] = channel;
2069                h[j * n + i] = channel;
2070            }
2071        }
2072        Self {
2073            arena: self.arena,
2074            v: d[0],
2075            g,
2076            h,
2077        }
2078    }
2079}
2080
2081/// Runtime-sized one-seed scalar for a Hessian-contracted third derivative.
2082#[derive(Clone, Copy, Debug)]
2083pub struct DynamicOneSeed<'arena> {
2084    /// Base value/gradient/Hessian channels.
2085    pub base: DynamicOrder2<'arena>,
2086    /// Nilpotent `epsilon` coefficient.
2087    pub eps: DynamicOrder2<'arena>,
2088}
2089
2090impl<'arena> DynamicOneSeed<'arena> {
2091    /// Seed one primary with the supplied contraction direction component.
2092    #[inline(always)]
2093    #[must_use]
2094    pub fn seed_direction(
2095        x: f64,
2096        axis: usize,
2097        u_axis: f64,
2098        dimension: usize,
2099        arena: &'arena DynamicJetArena,
2100    ) -> Self {
2101        Self {
2102            base: DynamicOrder2::variable(x, axis, dimension, arena),
2103            eps: DynamicOrder2::constant(u_axis, dimension, arena),
2104        }
2105    }
2106
2107    /// Row-major contracted-third matrix.
2108    #[inline(always)]
2109    #[must_use]
2110    pub fn contracted_third(&self) -> &[f64] {
2111        self.eps.h()
2112    }
2113}
2114
2115impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeed<'arena> {
2116    type Workspace = DynamicJetArena;
2117
2118    #[inline(always)]
2119    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2120        Self {
2121            base: DynamicOrder2::constant(c, dimension, arena),
2122            eps: DynamicOrder2::constant(0.0, dimension, arena),
2123        }
2124    }
2125
2126    #[inline(always)]
2127    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2128        Self {
2129            base: DynamicOrder2::variable(x, axis, dimension, arena),
2130            eps: DynamicOrder2::constant(0.0, dimension, arena),
2131        }
2132    }
2133
2134    #[inline(always)]
2135    fn dimension(&self) -> usize {
2136        self.base.dimension()
2137    }
2138
2139    #[inline(always)]
2140    fn value(&self) -> f64 {
2141        self.base.value()
2142    }
2143
2144    #[inline(always)]
2145    fn add(&self, o: &Self) -> Self {
2146        Self {
2147            base: self.base.add(&o.base),
2148            eps: self.eps.add(&o.eps),
2149        }
2150    }
2151
2152    #[inline(always)]
2153    fn sub(&self, o: &Self) -> Self {
2154        Self {
2155            base: self.base.sub(&o.base),
2156            eps: self.eps.sub(&o.eps),
2157        }
2158    }
2159
2160    #[inline(always)]
2161    fn mul(&self, o: &Self) -> Self {
2162        self.base.assert_compatible(&o.base);
2163        self.eps.assert_compatible(&o.eps);
2164        Self {
2165            base: self.base.mul(&o.base),
2166            eps: DynamicOrder2::from_channel_functions(
2167                self.base.v * o.eps.v + self.eps.v * o.base.v,
2168                self.dimension(),
2169                self.base.arena,
2170                |i| {
2171                    self.base.v * o.eps.g[i]
2172                        + self.base.g[i] * o.eps.v
2173                        + self.eps.v * o.base.g[i]
2174                        + self.eps.g[i] * o.base.v
2175                },
2176                |i, j| {
2177                    let ij = i * self.dimension() + j;
2178                    self.base.v * o.eps.h[ij]
2179                        + self.base.g[i] * o.eps.g[j]
2180                        + self.base.g[j] * o.eps.g[i]
2181                        + self.base.h[ij] * o.eps.v
2182                        + self.eps.v * o.base.h[ij]
2183                        + self.eps.g[i] * o.base.g[j]
2184                        + self.eps.g[j] * o.base.g[i]
2185                        + self.eps.h[ij] * o.base.v
2186                },
2187            ),
2188        }
2189    }
2190
2191    #[inline(always)]
2192    fn neg(&self) -> Self {
2193        Self {
2194            base: self.base.neg(),
2195            eps: self.eps.neg(),
2196        }
2197    }
2198
2199    #[inline(always)]
2200    fn scale(&self, s: f64) -> Self {
2201        Self {
2202            base: self.base.scale(s),
2203            eps: self.eps.scale(s),
2204        }
2205    }
2206
2207    #[inline(always)]
2208    fn compose_unary(&self, d: [f64; 5]) -> Self {
2209        let base = self.base.compose_unary(d);
2210        let dimension = self.dimension();
2211        let eps = DynamicOrder2::from_channel_functions(
2212            d[1] * self.eps.v,
2213            dimension,
2214            self.base.arena,
2215            |i| d[2] * self.base.g[i] * self.eps.v + d[1] * self.eps.g[i],
2216            |i, j| {
2217                let ij = i * dimension + j;
2218                d[1] * self.eps.h[ij]
2219                    + d[2]
2220                        * (self.base.g[i] * self.eps.g[j]
2221                            + self.base.g[j] * self.eps.g[i]
2222                            + self.base.h[ij] * self.eps.v)
2223                    + d[3] * self.base.g[i] * self.base.g[j] * self.eps.v
2224            },
2225        );
2226        Self { base, eps }
2227    }
2228}
2229
2230/// Reusable arena plus lane count for runtime directional batches.
2231///
2232/// A caller resets this once per row/chunk and evaluates every requested
2233/// third- or fourth-derivative contraction in one row-program pass. The bump
2234/// retains its largest chunk, so warmed rows do not return to the global
2235/// allocator.
2236#[derive(Debug)]
2237pub struct DynamicJetBatchWorkspace {
2238    arena: DynamicJetArena,
2239    lanes: usize,
2240}
2241
2242impl DynamicJetBatchWorkspace {
2243    /// Create a reusable workspace for `lanes` simultaneous directions.
2244    #[must_use]
2245    pub fn new(lanes: usize) -> Self {
2246        Self {
2247            arena: DynamicJetArena::new(),
2248            lanes,
2249        }
2250    }
2251
2252    /// Reclaim all jet storage and select the next evaluation's lane count.
2253    pub fn reset(&mut self, lanes: usize) {
2254        self.arena.reset();
2255        self.lanes = lanes;
2256    }
2257
2258    /// Bytes retained by the bump allocator after the largest evaluation.
2259    #[must_use]
2260    pub fn allocated_bytes(&self) -> usize {
2261        self.arena.allocated_bytes()
2262    }
2263
2264    /// Allocate a primary array in the same arena as every scalar channel.
2265    #[inline(always)]
2266    pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
2267        self.arena.alloc_slice_fill_with(len, fill)
2268    }
2269}
2270
2271/// Runtime-sized batch of one-seed contractions sharing one order-two base.
2272///
2273/// Lane `l` is algebraically identical to a standalone [`DynamicOneSeed`]
2274/// seeded by direction `u_l`, but the value/gradient/Hessian base is evaluated
2275/// once rather than once per direction. The output lane's epsilon Hessian is
2276/// `sum_c d³f/(dx_a dx_b dx_c) u_l[c]`.
2277#[derive(Clone, Copy, Debug)]
2278pub struct DynamicOneSeedBatch<'arena> {
2279    /// Shared value/gradient/Hessian channels.
2280    pub base: DynamicOrder2<'arena>,
2281    /// One nilpotent epsilon coefficient per contraction direction.
2282    eps: &'arena [DynamicOrder2<'arena>],
2283}
2284
2285impl<'arena> DynamicOneSeedBatch<'arena> {
2286    /// Seed one primary across every direction lane.
2287    #[inline(always)]
2288    #[must_use]
2289    pub fn seed_directions(
2290        x: f64,
2291        axis: usize,
2292        dimension: usize,
2293        workspace: &'arena DynamicJetBatchWorkspace,
2294        mut direction_at: impl FnMut(usize) -> f64,
2295    ) -> Self {
2296        let eps = workspace
2297            .arena
2298            .alloc_slice_fill_with(workspace.lanes, |lane| {
2299                DynamicOrder2::constant(direction_at(lane), dimension, &workspace.arena)
2300            });
2301        Self {
2302            base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2303            eps,
2304        }
2305    }
2306
2307    /// Number of simultaneous contraction directions.
2308    #[inline(always)]
2309    #[must_use]
2310    pub fn lanes(&self) -> usize {
2311        self.eps.len()
2312    }
2313
2314    /// Row-major contracted-third matrix for one direction lane.
2315    #[inline(always)]
2316    #[must_use]
2317    pub fn contracted_third(&self, lane: usize) -> &[f64] {
2318        self.eps[lane].h()
2319    }
2320
2321    #[inline(always)]
2322    fn assert_compatible(&self, other: &Self) {
2323        self.base.assert_compatible(&other.base);
2324        assert_eq!(
2325            self.eps.len(),
2326            other.eps.len(),
2327            "dynamic one-seed batch lane mismatch"
2328        );
2329    }
2330}
2331
2332impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeedBatch<'arena> {
2333    type Workspace = DynamicJetBatchWorkspace;
2334
2335    #[inline(always)]
2336    fn constant(c: f64, dimension: usize, workspace: &'arena DynamicJetBatchWorkspace) -> Self {
2337        let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2338            DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2339        });
2340        Self {
2341            base: DynamicOrder2::constant(c, dimension, &workspace.arena),
2342            eps,
2343        }
2344    }
2345
2346    #[inline(always)]
2347    fn variable(
2348        x: f64,
2349        axis: usize,
2350        dimension: usize,
2351        workspace: &'arena DynamicJetBatchWorkspace,
2352    ) -> Self {
2353        let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2354            DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2355        });
2356        Self {
2357            base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2358            eps,
2359        }
2360    }
2361
2362    #[inline(always)]
2363    fn dimension(&self) -> usize {
2364        self.base.dimension()
2365    }
2366
2367    #[inline(always)]
2368    fn value(&self) -> f64 {
2369        self.base.value()
2370    }
2371
2372    #[inline(always)]
2373    fn add(&self, other: &Self) -> Self {
2374        self.assert_compatible(other);
2375        let eps = self
2376            .base
2377            .arena
2378            .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].add(&other.eps[lane]));
2379        Self {
2380            base: self.base.add(&other.base),
2381            eps,
2382        }
2383    }
2384
2385    #[inline(always)]
2386    fn sub(&self, other: &Self) -> Self {
2387        self.assert_compatible(other);
2388        let eps = self
2389            .base
2390            .arena
2391            .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].sub(&other.eps[lane]));
2392        Self {
2393            base: self.base.sub(&other.base),
2394            eps,
2395        }
2396    }
2397
2398    #[inline(always)]
2399    fn mul(&self, other: &Self) -> Self {
2400        self.assert_compatible(other);
2401        let eps = self
2402            .base
2403            .arena
2404            .alloc_slice_fill_with(self.eps.len(), |lane| {
2405                self.base
2406                    .mul(&other.eps[lane])
2407                    .add(&self.eps[lane].mul(&other.base))
2408            });
2409        Self {
2410            base: self.base.mul(&other.base),
2411            eps,
2412        }
2413    }
2414
2415    #[inline(always)]
2416    fn neg(&self) -> Self {
2417        self.scale(-1.0)
2418    }
2419
2420    #[inline(always)]
2421    fn scale(&self, scale: f64) -> Self {
2422        let eps = self
2423            .base
2424            .arena
2425            .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].scale(scale));
2426        Self {
2427            base: self.base.scale(scale),
2428            eps,
2429        }
2430    }
2431
2432    #[inline(always)]
2433    fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
2434        let fprime = self.base.compose_unary([
2435            derivatives[1],
2436            derivatives[2],
2437            derivatives[3],
2438            derivatives[4],
2439            derivatives[4],
2440        ]);
2441        let eps = self
2442            .base
2443            .arena
2444            .alloc_slice_fill_with(self.eps.len(), |lane| fprime.mul(&self.eps[lane]));
2445        Self {
2446            base: self.base.compose_unary(derivatives),
2447            eps,
2448        }
2449    }
2450}
2451
2452/// Runtime-sized batch of two-seed contractions sharing one order-two base.
2453///
2454/// Lane `l` is algebraically identical to a standalone [`DynamicTwoSeed`]
2455/// seeded by the direction pair `(u_l, v_l)`, while the base row expression is
2456/// traversed once for the entire pair batch. Storage is `O(lanes * K^2)`; no
2457/// dense fourth-order tensor is formed.
2458#[derive(Clone, Copy, Debug)]
2459pub struct DynamicTwoSeedBatch<'arena> {
2460    /// Shared value/gradient/Hessian channels.
2461    pub base: DynamicOrder2<'arena>,
2462    eps: &'arena [DynamicOrder2<'arena>],
2463    del: &'arena [DynamicOrder2<'arena>],
2464    eps_del: &'arena [DynamicOrder2<'arena>],
2465}
2466
2467impl<'arena> DynamicTwoSeedBatch<'arena> {
2468    /// Seed one primary across every direction-pair lane.
2469    #[inline(always)]
2470    #[must_use]
2471    pub fn seed_direction_pairs(
2472        x: f64,
2473        axis: usize,
2474        dimension: usize,
2475        workspace: &'arena DynamicJetBatchWorkspace,
2476        mut direction_pair_at: impl FnMut(usize) -> (f64, f64),
2477    ) -> Self {
2478        let directions = workspace
2479            .arena
2480            .alloc_slice_fill_with(workspace.lanes, |lane| direction_pair_at(lane));
2481        let eps = workspace
2482            .arena
2483            .alloc_slice_fill_with(workspace.lanes, |lane| {
2484                DynamicOrder2::constant(directions[lane].0, dimension, &workspace.arena)
2485            });
2486        let del = workspace
2487            .arena
2488            .alloc_slice_fill_with(workspace.lanes, |lane| {
2489                DynamicOrder2::constant(directions[lane].1, dimension, &workspace.arena)
2490            });
2491        let eps_del = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2492            DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2493        });
2494        Self {
2495            base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2496            eps,
2497            del,
2498            eps_del,
2499        }
2500    }
2501
2502    /// Number of simultaneous contraction pairs.
2503    #[inline(always)]
2504    #[must_use]
2505    pub fn lanes(&self) -> usize {
2506        self.eps.len()
2507    }
2508
2509    /// Row-major contracted-fourth matrix for one direction-pair lane.
2510    #[inline(always)]
2511    #[must_use]
2512    pub fn contracted_fourth(&self, lane: usize) -> &[f64] {
2513        self.eps_del[lane].h()
2514    }
2515
2516    #[inline(always)]
2517    fn assert_compatible(&self, other: &Self) {
2518        self.base.assert_compatible(&other.base);
2519        assert_eq!(
2520            self.eps.len(),
2521            other.eps.len(),
2522            "dynamic two-seed batch lane mismatch"
2523        );
2524        assert_eq!(
2525            self.del.len(),
2526            self.eps.len(),
2527            "dynamic two-seed batch delta mismatch"
2528        );
2529        assert_eq!(
2530            self.eps_del.len(),
2531            self.eps.len(),
2532            "dynamic two-seed batch cross mismatch"
2533        );
2534    }
2535}
2536
2537impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeedBatch<'arena> {
2538    type Workspace = DynamicJetBatchWorkspace;
2539
2540    #[inline(always)]
2541    fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
2542        let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2543            DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2544        });
2545        Self {
2546            base: DynamicOrder2::constant(c, dimension, &workspace.arena),
2547            eps: zero,
2548            del: zero,
2549            eps_del: zero,
2550        }
2551    }
2552
2553    #[inline(always)]
2554    fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
2555        let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2556            DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2557        });
2558        Self {
2559            base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2560            eps: zero,
2561            del: zero,
2562            eps_del: zero,
2563        }
2564    }
2565
2566    #[inline(always)]
2567    fn dimension(&self) -> usize {
2568        self.base.dimension()
2569    }
2570
2571    #[inline(always)]
2572    fn value(&self) -> f64 {
2573        self.base.value()
2574    }
2575
2576    #[inline(always)]
2577    fn add(&self, other: &Self) -> Self {
2578        self.assert_compatible(other);
2579        let arena = self.base.arena;
2580        let eps =
2581            arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].add(&other.eps[lane]));
2582        let del =
2583            arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].add(&other.del[lane]));
2584        let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2585            self.eps_del[lane].add(&other.eps_del[lane])
2586        });
2587        Self {
2588            base: self.base.add(&other.base),
2589            eps,
2590            del,
2591            eps_del,
2592        }
2593    }
2594
2595    #[inline(always)]
2596    fn sub(&self, other: &Self) -> Self {
2597        self.assert_compatible(other);
2598        let arena = self.base.arena;
2599        let eps =
2600            arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].sub(&other.eps[lane]));
2601        let del =
2602            arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].sub(&other.del[lane]));
2603        let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2604            self.eps_del[lane].sub(&other.eps_del[lane])
2605        });
2606        Self {
2607            base: self.base.sub(&other.base),
2608            eps,
2609            del,
2610            eps_del,
2611        }
2612    }
2613
2614    #[inline(always)]
2615    fn mul(&self, other: &Self) -> Self {
2616        self.assert_compatible(other);
2617        let arena = self.base.arena;
2618        let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2619            self.base
2620                .mul(&other.eps[lane])
2621                .add(&self.eps[lane].mul(&other.base))
2622        });
2623        let del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2624            self.base
2625                .mul(&other.del[lane])
2626                .add(&self.del[lane].mul(&other.base))
2627        });
2628        let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2629            self.base
2630                .mul(&other.eps_del[lane])
2631                .add(&self.eps[lane].mul(&other.del[lane]))
2632                .add(&self.del[lane].mul(&other.eps[lane]))
2633                .add(&self.eps_del[lane].mul(&other.base))
2634        });
2635        Self {
2636            base: self.base.mul(&other.base),
2637            eps,
2638            del,
2639            eps_del,
2640        }
2641    }
2642
2643    #[inline(always)]
2644    fn neg(&self) -> Self {
2645        self.scale(-1.0)
2646    }
2647
2648    #[inline(always)]
2649    fn scale(&self, scale: f64) -> Self {
2650        let arena = self.base.arena;
2651        let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].scale(scale));
2652        let del = arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].scale(scale));
2653        let eps_del =
2654            arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps_del[lane].scale(scale));
2655        Self {
2656            base: self.base.scale(scale),
2657            eps,
2658            del,
2659            eps_del,
2660        }
2661    }
2662
2663    #[inline(always)]
2664    fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
2665        let arena = self.base.arena;
2666        let fprime = self.base.compose_unary([
2667            derivatives[1],
2668            derivatives[2],
2669            derivatives[3],
2670            derivatives[4],
2671            derivatives[4],
2672        ]);
2673        let fsecond = self.base.compose_unary([
2674            derivatives[2],
2675            derivatives[3],
2676            derivatives[4],
2677            derivatives[4],
2678            derivatives[4],
2679        ]);
2680        let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.eps[lane]));
2681        let del = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.del[lane]));
2682        let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2683            fsecond
2684                .mul(&self.eps[lane])
2685                .mul(&self.del[lane])
2686                .add(&fprime.mul(&self.eps_del[lane]))
2687        });
2688        Self {
2689            base: self.base.compose_unary(derivatives),
2690            eps,
2691            del,
2692            eps_del,
2693        }
2694    }
2695}
2696
2697/// Runtime-sized two-seed scalar for a Hessian-contracted fourth derivative.
2698#[derive(Clone, Copy, Debug)]
2699pub struct DynamicTwoSeed<'arena> {
2700    /// Base value/gradient/Hessian channels.
2701    pub base: DynamicOrder2<'arena>,
2702    /// Nilpotent `epsilon` coefficient.
2703    pub eps: DynamicOrder2<'arena>,
2704    /// Nilpotent `delta` coefficient.
2705    pub del: DynamicOrder2<'arena>,
2706    /// Mixed `epsilon delta` coefficient.
2707    pub eps_del: DynamicOrder2<'arena>,
2708}
2709
2710impl<'arena> DynamicTwoSeed<'arena> {
2711    /// Seed one primary with both contraction direction components.
2712    #[inline(always)]
2713    #[must_use]
2714    pub fn seed(
2715        x: f64,
2716        axis: usize,
2717        u_axis: f64,
2718        v_axis: f64,
2719        dimension: usize,
2720        arena: &'arena DynamicJetArena,
2721    ) -> Self {
2722        Self {
2723            base: DynamicOrder2::variable(x, axis, dimension, arena),
2724            eps: DynamicOrder2::constant(u_axis, dimension, arena),
2725            del: DynamicOrder2::constant(v_axis, dimension, arena),
2726            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2727        }
2728    }
2729
2730    /// Row-major contracted-fourth matrix.
2731    #[inline(always)]
2732    #[must_use]
2733    pub fn contracted_fourth(&self) -> &[f64] {
2734        self.eps_del.h()
2735    }
2736}
2737
2738impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeed<'arena> {
2739    type Workspace = DynamicJetArena;
2740
2741    #[inline(always)]
2742    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2743        Self {
2744            base: DynamicOrder2::constant(c, dimension, arena),
2745            eps: DynamicOrder2::constant(0.0, dimension, arena),
2746            del: DynamicOrder2::constant(0.0, dimension, arena),
2747            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2748        }
2749    }
2750
2751    #[inline(always)]
2752    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2753        Self {
2754            base: DynamicOrder2::variable(x, axis, dimension, arena),
2755            eps: DynamicOrder2::constant(0.0, dimension, arena),
2756            del: DynamicOrder2::constant(0.0, dimension, arena),
2757            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2758        }
2759    }
2760
2761    #[inline(always)]
2762    fn dimension(&self) -> usize {
2763        self.base.dimension()
2764    }
2765
2766    #[inline(always)]
2767    fn value(&self) -> f64 {
2768        self.base.value()
2769    }
2770
2771    #[inline(always)]
2772    fn add(&self, o: &Self) -> Self {
2773        Self {
2774            base: self.base.add(&o.base),
2775            eps: self.eps.add(&o.eps),
2776            del: self.del.add(&o.del),
2777            eps_del: self.eps_del.add(&o.eps_del),
2778        }
2779    }
2780
2781    #[inline(always)]
2782    fn sub(&self, o: &Self) -> Self {
2783        Self {
2784            base: self.base.sub(&o.base),
2785            eps: self.eps.sub(&o.eps),
2786            del: self.del.sub(&o.del),
2787            eps_del: self.eps_del.sub(&o.eps_del),
2788        }
2789    }
2790
2791    #[inline(always)]
2792    fn mul(&self, o: &Self) -> Self {
2793        let base = self.base.mul(&o.base);
2794        let eps = self.base.mul(&o.eps).add(&self.eps.mul(&o.base));
2795        let del = self.base.mul(&o.del).add(&self.del.mul(&o.base));
2796        let eps_del = self
2797            .base
2798            .mul(&o.eps_del)
2799            .add(&self.eps.mul(&o.del))
2800            .add(&self.del.mul(&o.eps))
2801            .add(&self.eps_del.mul(&o.base));
2802        Self {
2803            base,
2804            eps,
2805            del,
2806            eps_del,
2807        }
2808    }
2809
2810    #[inline(always)]
2811    fn neg(&self) -> Self {
2812        Self {
2813            base: self.base.neg(),
2814            eps: self.eps.neg(),
2815            del: self.del.neg(),
2816            eps_del: self.eps_del.neg(),
2817        }
2818    }
2819
2820    #[inline(always)]
2821    fn scale(&self, s: f64) -> Self {
2822        Self {
2823            base: self.base.scale(s),
2824            eps: self.eps.scale(s),
2825            del: self.del.scale(s),
2826            eps_del: self.eps_del.scale(s),
2827        }
2828    }
2829
2830    #[inline(always)]
2831    fn compose_unary(&self, d: [f64; 5]) -> Self {
2832        let base = self.base.compose_unary(d);
2833        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
2834        let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]);
2835        let eps = fprime.mul(&self.eps);
2836        let del = fprime.mul(&self.del);
2837        let eps_del = fsecond
2838            .mul(&self.eps)
2839            .mul(&self.del)
2840            .add(&fprime.mul(&self.eps_del));
2841        Self {
2842            base,
2843            eps,
2844            del,
2845            eps_del,
2846        }
2847    }
2848}
2849
2850// ── Order2<K> ergonomic operator overloads (doc §A.1) ───────────────────
2851//
2852// The dispersion-family row NLLs are written with `+`/`-`/`*` operators over
2853// the primaries (mirroring how they read as `Tower4` expressions). These
2854// delegate channel-for-channel to the inner `Tower2` arithmetic (which has
2855// `Add`/`Mul`; `Sub`/`Neg` are expressed as `+ (-1)·rhs` exactly as the
2856// `crate::nested_dual::JetField::sub` / `crate::nested_dual::JetField::neg` impls do), so an `Order2` expression is
2857// bit-identical to the same `Tower4` expression's order-≤2 channels.
2858
2859impl<const K: usize> std::ops::Add for Order2<K> {
2860    type Output = Self;
2861    #[inline]
2862    fn add(self, o: Self) -> Self {
2863        Order2(self.0 + o.0)
2864    }
2865}
2866
2867impl<const K: usize> std::ops::Add<f64> for Order2<K> {
2868    type Output = Self;
2869    #[inline]
2870    fn add(self, c: f64) -> Self {
2871        Order2(self.0 + c)
2872    }
2873}
2874
2875impl<const K: usize> std::ops::Sub for Order2<K> {
2876    type Output = Self;
2877    #[inline]
2878    fn sub(self, o: Self) -> Self {
2879        Order2(self.0 + o.0.scale(-1.0))
2880    }
2881}
2882
2883impl<const K: usize> std::ops::Sub<f64> for Order2<K> {
2884    type Output = Self;
2885    #[inline]
2886    fn sub(self, c: f64) -> Self {
2887        Order2(self.0 + (-c))
2888    }
2889}
2890
2891impl<const K: usize> std::ops::Mul for Order2<K> {
2892    type Output = Self;
2893    #[inline]
2894    fn mul(self, o: Self) -> Self {
2895        Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
2896    }
2897}
2898
2899impl<const K: usize> std::ops::Mul<f64> for Order2<K> {
2900    type Output = Self;
2901    #[inline]
2902    fn mul(self, c: f64) -> Self {
2903        Order2(self.0.scale(c))
2904    }
2905}
2906
2907impl<const K: usize> std::ops::Neg for Order2<K> {
2908    type Output = Self;
2909    #[inline]
2910    fn neg(self) -> Self {
2911        Order2(self.0.scale(-1.0))
2912    }
2913}
2914
2915/// Filtered Hensel lift of a SCALAR implicit state `a(θ)` defined by the
2916/// constraint `F(a, θ) = 0`, evaluated in ANY [`JetScalar`] algebra `S` (doc
2917/// §11, "A generic implicit-lift operator for every production scalar").
2918///
2919/// This is the perf-respecting alternative to lifting through a dense
2920/// `Tower4<K+1>` (which carries the implicit variable as an extra dense axis):
2921/// the state `a` lives directly in the consumer's own `K`-primary algebra
2922/// `S` — `Order2<K>` for value/gradient/Hessian, `Tower4<K>` for the full
2923/// `t3`/`t4` — never paying for an extra variable.
2924///
2925/// **Method.** Fixed-Jacobian Newton in the nilpotent algebra. By the
2926/// filtered-lift theorem (doc §11.1), if `F_a := ∂F/∂a(a₀, θ₀)` is the primal
2927/// Jacobian at the base point and `inv_fa = 1/F_a`, then the iteration
2928/// `A ← A − inv_fa · F(A, θ)` raises the filtration degree of the residual by
2929/// at least one per step: each step kills exactly one graded layer. Starting
2930/// from `A = const(a₀)` (whose residual lies in `F¹` because `θ − θ₀ ∈ 𝔫`),
2931/// `iters` equal to the algebra's nilpotency order returns the *exact* lifted
2932/// jet (`Order2`: 2, `OneSeed`: 3, `Tower4`/`TwoSeed`: 4). The value channel of
2933/// `A` never moves — `F(A, θ).value() = F(a₀, θ₀) = 0` at the certified root —
2934/// so a caller may precompute every primitive's derivative stack at the fixed
2935/// base index once and let the cheap polynomial composition repeat per step.
2936///
2937/// `f` evaluates the constraint `F(a, θ)` in `S` (capturing the seeded
2938/// parameter jets `θ`); `a0` is the certified scalar root `F(a₀, θ₀) ≈ 0`.
2939pub fn filtered_implicit_solve_scalar<const K: usize, S: JetScalar<K>>(
2940    a0: f64,
2941    inv_fa: f64,
2942    iters: usize,
2943    f: impl Fn(&S) -> S,
2944) -> S {
2945    let mut a = S::constant(a0);
2946    for _ in 0..iters {
2947        let residual = f(&a);
2948        a = a.sub(&residual.scale(inv_fa));
2949    }
2950    a
2951}
2952
2953// ── PatternedOrder2<P, K, H>: compile-time sparse Hessian ───────────────
2954
2955/// Runtime-dimension form of [`filtered_implicit_solve_scalar`].
2956///
2957/// The state is lifted directly in the caller's [`RuntimeJetScalar`] algebra,
2958/// so one row expression can serve a runtime-sized value/gradient/Hessian
2959/// evaluation (`DynamicOrder2`), a Hessian-contracted third derivative
2960/// (`DynamicOneSeed`), and a Hessian-contracted fourth derivative
2961/// (`DynamicTwoSeed`). `iters` is the algebra's nilpotency order: two, three,
2962/// and four respectively. No implicit axis is appended to the derivative
2963/// storage.
2964pub fn filtered_implicit_solve_runtime_scalar<'arena, S: RuntimeJetScalar<'arena>>(
2965    a0: f64,
2966    inv_fa: f64,
2967    iters: usize,
2968    dimension: usize,
2969    workspace: &'arena S::Workspace,
2970    f: impl Fn(&S) -> S,
2971) -> S {
2972    let mut a = S::constant(a0, dimension, workspace);
2973    for _ in 0..iters {
2974        let residual = f(&a);
2975        a = a.sub(&residual.scale(inv_fa));
2976    }
2977    a
2978}
2979
2980/// Compile-time upper-triangle Hessian pattern for [`PatternedOrder2`].
2981///
2982/// `PAIRS` contains exactly the `(row, column)` channels a row program can
2983/// produce, with `row <= column`.  Pointwise Taylor algebra never couples one
2984/// Hessian pair through another: output `H[i,j]` depends only on the two input
2985/// `H[i,j]` channels and their `g[i]`/`g[j]` channels.  It is therefore exact to
2986/// omit structurally impossible pairs rather than multiplying their zeros.
2987pub trait HessianPattern<const K: usize, const H: usize> {
2988    const PAIRS: [(usize, usize); H];
2989    const PAIR_BITS: [[u128; K]; K];
2990}
2991
2992/// Build the symmetric axis-pair → patterned-slot lookup used by dependency
2993/// propagation in [`PatternedOrder2`].
2994pub const fn hessian_pair_bits<const K: usize, const H: usize>(
2995    pairs: [(usize, usize); H],
2996) -> [[u128; K]; K] {
2997    let mut table = [[0u128; K]; K];
2998    let mut slot = 0;
2999    while slot < H {
3000        let (i, j) = pairs[slot];
3001        let bit = 1u128 << slot;
3002        table[i][j] = bit;
3003        table[j][i] = bit;
3004        slot += 1;
3005    }
3006    table
3007}
3008
3009/// Exact order-two jet with a dense gradient and a compile-time patterned
3010/// upper-triangle Hessian.
3011///
3012/// This carries `1 + K + H` scalars instead of `1 + K + K²`.  Its arithmetic is
3013/// the same Leibniz/Faà-di-Bruno algebra as [`Order2`], evaluated only for the
3014/// Hessian pairs declared by `P`.  A family row NLL remains written once over
3015/// [`JetScalar`]; the pattern is an execution schedule, not a second derivative
3016/// formula.
3017#[derive(Debug)]
3018pub struct PatternedOrder2<P, const K: usize, const H: usize> {
3019    v: f64,
3020    g: [f64; K],
3021    h: [f64; H],
3022    gradient_mask: u128,
3023    hessian_mask: u128,
3024    pattern: std::marker::PhantomData<fn() -> P>,
3025}
3026
3027impl<P, const K: usize, const H: usize> Copy for PatternedOrder2<P, K, H> {}
3028
3029impl<P, const K: usize, const H: usize> Clone for PatternedOrder2<P, K, H> {
3030    fn clone(&self) -> Self {
3031        *self
3032    }
3033}
3034
3035impl<P, const K: usize, const H: usize> PatternedOrder2<P, K, H>
3036where
3037    P: HessianPattern<K, H>,
3038{
3039    #[inline]
3040    #[must_use]
3041    pub fn g(&self) -> [f64; K] {
3042        self.g
3043    }
3044
3045    /// Expand the patterned upper triangle into the dense symmetric shape
3046    /// required by the existing `RowKernel` interface. Missing pairs are exact
3047    /// structural zeros.
3048    #[inline]
3049    #[must_use]
3050    pub fn h(&self) -> [[f64; K]; K] {
3051        let mut dense = [[0.0; K]; K];
3052        for (slot, &(i, j)) in P::PAIRS.iter().enumerate() {
3053            dense[i][j] = self.h[slot];
3054            dense[j][i] = self.h[slot];
3055        }
3056        dense
3057    }
3058
3059    #[inline]
3060    fn pair_mask_between(left: u128, right: u128) -> u128 {
3061        let mut result = 0u128;
3062        let mut left_axes = left;
3063        while left_axes != 0 {
3064            let i = left_axes.trailing_zeros() as usize;
3065            left_axes &= left_axes - 1;
3066            let mut right_axes = right;
3067            while right_axes != 0 {
3068                let j = right_axes.trailing_zeros() as usize;
3069                right_axes &= right_axes - 1;
3070                result |= P::PAIR_BITS[i][j];
3071            }
3072        }
3073        result
3074    }
3075}
3076
3077impl<P, const K: usize, const H: usize> JetScalar<K> for PatternedOrder2<P, K, H>
3078where
3079    P: HessianPattern<K, H>,
3080{
3081    #[inline]
3082    fn constant(c: f64) -> Self {
3083        Self {
3084            v: c,
3085            g: [0.0; K],
3086            h: [0.0; H],
3087            gradient_mask: 0,
3088            hessian_mask: 0,
3089            pattern: std::marker::PhantomData,
3090        }
3091    }
3092
3093    #[inline]
3094    fn variable(x: f64, axis: usize) -> Self {
3095        let mut out = Self::constant(x);
3096        if axis < K {
3097            out.g[axis] = 1.0;
3098            out.gradient_mask = 1u128 << axis;
3099        }
3100        out
3101    }
3102}
3103
3104impl<P, const K: usize, const H: usize> crate::nested_dual::JetField for PatternedOrder2<P, K, H>
3105where
3106    P: HessianPattern<K, H>,
3107{
3108    #[inline]
3109    fn value(&self) -> f64 {
3110        self.v
3111    }
3112
3113    #[inline]
3114    fn add(&self, other: &Self) -> Self {
3115        let mut out = Self::constant(self.v + other.v);
3116        out.gradient_mask = self.gradient_mask | other.gradient_mask;
3117        let mut gradient_mask = out.gradient_mask;
3118        while gradient_mask != 0 {
3119            let i = gradient_mask.trailing_zeros() as usize;
3120            gradient_mask &= gradient_mask - 1;
3121            out.g[i] = self.g[i] + other.g[i];
3122        }
3123        out.hessian_mask = self.hessian_mask | other.hessian_mask;
3124        let mut hessian_mask = out.hessian_mask;
3125        while hessian_mask != 0 {
3126            let slot = hessian_mask.trailing_zeros() as usize;
3127            hessian_mask &= hessian_mask - 1;
3128            out.h[slot] = self.h[slot] + other.h[slot];
3129        }
3130        out
3131    }
3132
3133    #[inline]
3134    fn sub(&self, other: &Self) -> Self {
3135        let mut out = Self::constant(self.v - other.v);
3136        out.gradient_mask = self.gradient_mask | other.gradient_mask;
3137        let mut gradient_mask = out.gradient_mask;
3138        while gradient_mask != 0 {
3139            let i = gradient_mask.trailing_zeros() as usize;
3140            gradient_mask &= gradient_mask - 1;
3141            out.g[i] = self.g[i] - other.g[i];
3142        }
3143        out.hessian_mask = self.hessian_mask | other.hessian_mask;
3144        let mut hessian_mask = out.hessian_mask;
3145        while hessian_mask != 0 {
3146            let slot = hessian_mask.trailing_zeros() as usize;
3147            hessian_mask &= hessian_mask - 1;
3148            out.h[slot] = self.h[slot] - other.h[slot];
3149        }
3150        out
3151    }
3152
3153    #[inline]
3154    fn mul(&self, other: &Self) -> Self {
3155        let mut out = Self::constant(self.v * other.v);
3156        out.gradient_mask = self.gradient_mask | other.gradient_mask;
3157        let mut gradient_mask = out.gradient_mask;
3158        while gradient_mask != 0 {
3159            let i = gradient_mask.trailing_zeros() as usize;
3160            gradient_mask &= gradient_mask - 1;
3161            out.g[i] = self.v * other.g[i] + self.g[i] * other.v;
3162        }
3163        out.hessian_mask = self.hessian_mask
3164            | other.hessian_mask
3165            | Self::pair_mask_between(self.gradient_mask, other.gradient_mask);
3166        let mut hessian_mask = out.hessian_mask;
3167        while hessian_mask != 0 {
3168            let slot = hessian_mask.trailing_zeros() as usize;
3169            hessian_mask &= hessian_mask - 1;
3170            let (i, j) = P::PAIRS[slot];
3171            out.h[slot] = self.v * other.h[slot]
3172                + self.g[i] * other.g[j]
3173                + self.g[j] * other.g[i]
3174                + self.h[slot] * other.v;
3175        }
3176        out
3177    }
3178
3179    #[inline]
3180    fn neg(&self) -> Self {
3181        self.scale(-1.0)
3182    }
3183
3184    #[inline]
3185    fn scale(&self, scale: f64) -> Self {
3186        let mut out = Self::constant(self.v * scale);
3187        out.gradient_mask = self.gradient_mask;
3188        let mut gradient_mask = out.gradient_mask;
3189        while gradient_mask != 0 {
3190            let i = gradient_mask.trailing_zeros() as usize;
3191            gradient_mask &= gradient_mask - 1;
3192            out.g[i] = self.g[i] * scale;
3193        }
3194        out.hessian_mask = self.hessian_mask;
3195        let mut hessian_mask = out.hessian_mask;
3196        while hessian_mask != 0 {
3197            let slot = hessian_mask.trailing_zeros() as usize;
3198            hessian_mask &= hessian_mask - 1;
3199            out.h[slot] = self.h[slot] * scale;
3200        }
3201        out
3202    }
3203
3204    #[inline]
3205    fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
3206        let mut out = Self::constant(derivatives[0]);
3207        out.gradient_mask = self.gradient_mask;
3208        let mut gradient_mask = out.gradient_mask;
3209        while gradient_mask != 0 {
3210            let i = gradient_mask.trailing_zeros() as usize;
3211            gradient_mask &= gradient_mask - 1;
3212            out.g[i] = derivatives[1] * self.g[i];
3213        }
3214        out.hessian_mask =
3215            self.hessian_mask | Self::pair_mask_between(self.gradient_mask, self.gradient_mask);
3216        let mut hessian_mask = out.hessian_mask;
3217        while hessian_mask != 0 {
3218            let slot = hessian_mask.trailing_zeros() as usize;
3219            hessian_mask &= hessian_mask - 1;
3220            let (i, j) = P::PAIRS[slot];
3221            out.h[slot] = derivatives[2] * self.g[i] * self.g[j] + derivatives[1] * self.h[slot];
3222        }
3223        out
3224    }
3225}
3226
3227// ── Order2<K>: value / gradient / Hessian (doc §A.1) ────────────────────
3228
3229/// Truncated SECOND-order scalar: value `v`, gradient `g_a`, Hessian `H_{ab}`.
3230///
3231/// This is a thin newtype over the production [`crate::jet_tower::Tower2`], so
3232/// its `(v, g, H)` channels are obtained by the SAME formulas — and are
3233/// therefore bit-identical to both [`crate::jet_tower::Tower2`] and the order-≤2
3234/// channels of a full [`crate::jet_tower::Tower4`] (doc §A.1, "Bit-identity with
3235/// the full tower"). The wrapper exists only to satisfy the generic
3236/// [`JetScalar`] interface (the `compose_unary` / `add` / `sub` / `neg` /
3237/// `recip` the trait demands, which `Tower2` does not expose by that shape) —
3238/// every channel is delegated to `Tower2` arithmetic unchanged.
3239#[derive(Clone, Copy, Debug)]
3240pub struct Order2<const K: usize>(pub crate::jet_tower::Tower2<K>);
3241
3242impl<const K: usize> Order2<K> {
3243    /// Read the gradient channel `g_a = ∂ℓ/∂p_a`.
3244    #[inline]
3245    #[must_use]
3246    pub fn g(&self) -> &[f64; K] {
3247        &self.0.g
3248    }
3249
3250    /// Read the Hessian channel.
3251    #[inline]
3252    #[must_use]
3253    pub fn h(&self) -> &[[f64; K]; K] {
3254        &self.0.h
3255    }
3256
3257    /// Consume the jet and move out its value, gradient, and Hessian channels.
3258    #[inline]
3259    #[must_use]
3260    pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
3261        let crate::jet_tower::Tower2 { v, g, h } = self.0;
3262        (v, g, h)
3263    }
3264}
3265
3266impl<const K: usize> JetScalar<K> for Order2<K> {
3267    fn constant(c: f64) -> Self {
3268        Order2(crate::jet_tower::Tower2::constant(c))
3269    }
3270    fn variable(x: f64, axis: usize) -> Self {
3271        Order2(crate::jet_tower::Tower2::variable(x, axis))
3272    }
3273
3274    #[inline(always)]
3275    fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
3276        inputs: &[Self],
3277        coefficients: &C,
3278    ) -> Self {
3279        assert_eq!(inputs.len(), coefficients.dimension());
3280        let input_dimension = inputs.len();
3281        assert!(input_dimension <= K);
3282        let mut values = [0.0; K];
3283        for axis in 0..input_dimension {
3284            values[axis] = inputs[axis].0.v;
3285        }
3286        let mut projected = [0.0; K];
3287        coefficients.multiply(
3288            &values[..input_dimension],
3289            &mut projected[..input_dimension],
3290        );
3291
3292        let mut out = crate::jet_tower::Tower2::zero();
3293        for axis in 0..input_dimension {
3294            out.v += values[axis] * projected[axis];
3295        }
3296        for primary in 0..K {
3297            let mut channel = 0.0;
3298            for axis in 0..input_dimension {
3299                channel += projected[axis] * inputs[axis].0.g[primary];
3300            }
3301            out.g[primary] = 2.0 * channel;
3302        }
3303        let mut input_gradient = [0.0; K];
3304        let mut projected_gradient = [0.0; K];
3305        for primary_b in 0..K {
3306            for row in 0..input_dimension {
3307                input_gradient[row] = inputs[row].0.g[primary_b];
3308            }
3309            coefficients.multiply(
3310                &input_gradient[..input_dimension],
3311                &mut projected_gradient[..input_dimension],
3312            );
3313            for primary_a in 0..=primary_b {
3314                let mut inherited = 0.0;
3315                let mut curvature = 0.0;
3316                for row in 0..input_dimension {
3317                    inherited += projected[row] * inputs[row].0.h[primary_a][primary_b];
3318                    curvature += inputs[row].0.g[primary_a] * projected_gradient[row];
3319                }
3320                let channel = 2.0 * (inherited + curvature);
3321                out.h[primary_a][primary_b] = channel;
3322                out.h[primary_b][primary_a] = channel;
3323            }
3324        }
3325        Order2(out)
3326    }
3327
3328    #[inline(always)]
3329    fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
3330        assert_eq!(inputs.len(), weights.len());
3331        let mut out = crate::jet_tower::Tower2::zero();
3332        for (input, &weight) in inputs.iter().zip(weights) {
3333            out.v += input.0.v * weight;
3334        }
3335        for primary in 0..K {
3336            for (input, &weight) in inputs.iter().zip(weights) {
3337                out.g[primary] += input.0.g[primary] * weight;
3338            }
3339            for other in primary..K {
3340                for (input, &weight) in inputs.iter().zip(weights) {
3341                    out.h[primary][other] += input.0.h[primary][other] * weight;
3342                }
3343                out.h[other][primary] = out.h[primary][other];
3344            }
3345        }
3346        Order2(out)
3347    }
3348
3349    #[inline(always)]
3350    fn add_constant(&self, constant: f64) -> Self {
3351        let mut out = *self;
3352        out.0.v += constant;
3353        out
3354    }
3355
3356    #[inline(always)]
3357    fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
3358        let mut out = crate::jet_tower::Tower2::zero();
3359        out.v = self.0.v * right.0.v + addend.0.v;
3360        for primary in 0..K {
3361            out.g[primary] =
3362                self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v + addend.0.g[primary];
3363            for other in primary..K {
3364                let channel = self.0.v * right.0.h[primary][other]
3365                    + self.0.g[primary] * right.0.g[other]
3366                    + self.0.g[other] * right.0.g[primary]
3367                    + self.0.h[primary][other] * right.0.v
3368                    + addend.0.h[primary][other];
3369                out.h[primary][other] = channel;
3370                out.h[other][primary] = channel;
3371            }
3372        }
3373        Order2(out)
3374    }
3375
3376    #[inline(always)]
3377    fn product(&self, right: &Self) -> Self {
3378        let mut out = crate::jet_tower::Tower2::zero();
3379        out.v = self.0.v * right.0.v;
3380        for primary in 0..K {
3381            out.g[primary] = self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v;
3382            for other in primary..K {
3383                let channel = self.0.v * right.0.h[primary][other]
3384                    + self.0.g[primary] * right.0.g[other]
3385                    + self.0.g[other] * right.0.g[primary]
3386                    + self.0.h[primary][other] * right.0.v;
3387                out.h[primary][other] = channel;
3388                out.h[other][primary] = channel;
3389            }
3390        }
3391        Order2(out)
3392    }
3393
3394    #[inline(always)]
3395    fn affine_compose(
3396        &self,
3397        input_scale: f64,
3398        input_shift: f64,
3399        derivative_stack: [f64; 5],
3400    ) -> Self {
3401        assert!(input_shift.is_finite(), "affine input shift must be finite");
3402        let first = derivative_stack[1] * input_scale;
3403        let second = derivative_stack[2] * input_scale * input_scale;
3404        let mut out = crate::jet_tower::Tower2::zero();
3405        out.v = derivative_stack[0];
3406        for primary in 0..K {
3407            out.g[primary] = first * self.0.g[primary];
3408            for other in primary..K {
3409                let channel =
3410                    first * self.0.h[primary][other] + second * self.0.g[primary] * self.0.g[other];
3411                out.h[primary][other] = channel;
3412                out.h[other][primary] = channel;
3413            }
3414        }
3415        Order2(out)
3416    }
3417
3418    #[inline(always)]
3419    fn affine_composed_sum(
3420        inputs: &[Self],
3421        input_scales: &[f64],
3422        derivative_stacks: &[[f64; 5]],
3423    ) -> Self {
3424        assert_eq!(inputs.len(), input_scales.len());
3425        assert_eq!(inputs.len(), derivative_stacks.len());
3426        let mut out = crate::jet_tower::Tower2::zero();
3427        for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
3428        {
3429            let first = stack[1] * input_scale;
3430            let second = stack[2] * input_scale * input_scale;
3431            out.v += stack[0];
3432            for primary in 0..K {
3433                out.g[primary] += first * input.0.g[primary];
3434                for other in primary..K {
3435                    out.h[primary][other] += first * input.0.h[primary][other]
3436                        + second * input.0.g[primary] * input.0.g[other];
3437                }
3438            }
3439        }
3440        for primary in 0..K {
3441            for other in primary + 1..K {
3442                out.h[other][primary] = out.h[primary][other];
3443            }
3444        }
3445        Order2(out)
3446    }
3447
3448    #[inline(always)]
3449    fn shared_multiply_add_affine_composed_sum<const N: usize>(
3450        lefts: &[&Self; N],
3451        right: &Self,
3452        addend: &Self,
3453        addend_scales: &[f64; N],
3454        input_scales: &[f64; N],
3455        derivative_stacks: &[[f64; 5]; N],
3456    ) -> Self {
3457        let (representatives, term_sources, source_count) =
3458            canonical_shared_source_schedule::<N>(|term, representative| {
3459                std::ptr::eq(lefts[term], lefts[representative])
3460                    && addend_scales[term] == addend_scales[representative]
3461            });
3462        let (value, source_derivatives) =
3463            aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
3464        let mut source_gradients = [[0.0; K]; N];
3465        let mut out = crate::jet_tower::Tower2::zero();
3466        out.v = value;
3467        let mut right_first = 0.0;
3468        let mut addend_first = 0.0;
3469        for source in 0..source_count {
3470            let term = representatives[source];
3471            let first = source_derivatives[source][1];
3472            right_first += first * lefts[term].0.v;
3473            addend_first += first * addend_scales[term];
3474            for primary in 0..K {
3475                let product_gradient =
3476                    lefts[term].0.v * right.0.g[primary] + lefts[term].0.g[primary] * right.0.v;
3477                let inner_gradient = if addend_scales[term] == 0.0 {
3478                    product_gradient
3479                } else if addend_scales[term] == 1.0 {
3480                    product_gradient + addend.0.g[primary]
3481                } else {
3482                    product_gradient + addend_scales[term] * addend.0.g[primary]
3483                };
3484                source_gradients[source][primary] = inner_gradient;
3485                out.g[primary] += first * lefts[term].0.g[primary] * right.0.v;
3486            }
3487        }
3488        if N != 0 {
3489            for primary in 0..K {
3490                out.g[primary] += right_first * right.0.g[primary];
3491            }
3492        }
3493        let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
3494        if addend_live {
3495            for primary in 0..K {
3496                out.g[primary] += addend_first * addend.0.g[primary];
3497            }
3498        }
3499        for primary in 0..K {
3500            for other in primary..K {
3501                let mut channel = if N == 0 {
3502                    0.0
3503                } else {
3504                    right_first * right.0.h[primary][other]
3505                };
3506                if addend_live {
3507                    channel += addend_first * addend.0.h[primary][other];
3508                }
3509                for source in 0..source_count {
3510                    let term = representatives[source];
3511                    let local_product_hessian = lefts[term].0.g[primary] * right.0.g[other]
3512                        + lefts[term].0.g[other] * right.0.g[primary]
3513                        + lefts[term].0.h[primary][other] * right.0.v;
3514                    channel += source_derivatives[source][1] * local_product_hessian
3515                        + source_derivatives[source][2]
3516                            * source_gradients[source][primary]
3517                            * source_gradients[source][other];
3518                }
3519                out.h[primary][other] = channel;
3520                out.h[other][primary] = channel;
3521            }
3522        }
3523        Order2(out)
3524    }
3525
3526    #[inline(always)]
3527    fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
3528        assert_eq!(inputs.len(), derivative_stacks.len());
3529        let mut out = crate::jet_tower::Tower2::zero();
3530        for (input, stack) in inputs.iter().zip(derivative_stacks) {
3531            out.v += stack[0];
3532            for primary in 0..K {
3533                out.g[primary] += stack[1] * input.0.g[primary];
3534                for other in primary..K {
3535                    out.h[primary][other] += stack[1] * input.0.h[primary][other]
3536                        + stack[2] * input.0.g[primary] * input.0.g[other];
3537                }
3538            }
3539        }
3540        for primary in 0..K {
3541            for other in primary + 1..K {
3542                out.h[other][primary] = out.h[primary][other];
3543            }
3544        }
3545        Order2(out)
3546    }
3547}
3548
3549impl<const K: usize> crate::nested_dual::JetField for Order2<K> {
3550    fn value(&self) -> f64 {
3551        self.0.v
3552    }
3553    fn add(&self, o: &Self) -> Self {
3554        Order2(self.0 + o.0)
3555    }
3556    fn sub(&self, o: &Self) -> Self {
3557        // Tower2 has no Sub op; subtract by adding the negation, matching
3558        // Tower4::sub (self + o.scale(-1.0)).
3559        Order2(self.0 + o.0.scale(-1.0))
3560    }
3561    fn mul(&self, o: &Self) -> Self {
3562        Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
3563    }
3564    fn neg(&self) -> Self {
3565        Order2(self.0.scale(-1.0))
3566    }
3567    fn scale(&self, s: f64) -> Self {
3568        Order2(self.0.scale(s))
3569    }
3570    fn compose_unary(&self, d: [f64; 5]) -> Self {
3571        // Order-≤2 reads only [f, f', f''] of the stack.
3572        Order2(self.0.compose_unary([d[0], d[1], d[2]]))
3573    }
3574}
3575
3576/// Static lowering target for a sum of composed, low-dimensional row atoms.
3577///
3578/// A row likelihood often depends on a large global primary vector only through
3579/// a few small independent indices.  Evaluating the whole expression in an
3580/// `Order2<K>` then pays dense `K²` arithmetic at every intermediate; carrying
3581/// runtime dependency masks replaces that arithmetic with branches and bit
3582/// scans.  This accumulator provides the ahead-of-time alternative: evaluate
3583/// each index once in its natural local dimension `N`, then scatter the exact
3584/// second-order composition into the global `K` channels through a fixed axis
3585/// map.  The family still owns only its scalar index expression and certified
3586/// unary derivative stack; this type owns the universal chain rule.
3587///
3588/// For an atom `q(x_local)` and outer stack `[f(q), f'(q), f''(q)]`, the lowered
3589/// channels are
3590///
3591/// ```text
3592/// g[a_i]       += f' q_i
3593/// H[a_i, a_j] += f' q_ij + f'' q_i q_j.
3594/// ```
3595///
3596/// Only the local upper triangle is evaluated, then mirrored into the global
3597/// symmetric output.  With literal `axes` and fixed `N`, LLVM unrolls this into
3598/// straight-line arithmetic: no dependency masks, sparse-pair lookups, or jet
3599/// temporaries survive into the generated schedule.
3600#[derive(Clone, Copy, Debug)]
3601pub struct MappedOrder2Accumulator<const K: usize> {
3602    value: f64,
3603    gradient: [f64; K],
3604    hessian: [[f64; K]; K],
3605}
3606
3607/// Compile-time-symbolic value/gradient/Hessian of one local row atom.
3608///
3609/// The Hessian stores only its upper triangle, in row-major triangular order.
3610/// Instances are emitted by `gam-row-macros`; unlike a runtime forward jet,
3611/// they contain only final live channels and therefore introduce no seeded
3612/// identity arrays, dependency masks, or zero arithmetic.
3613#[derive(Clone, Copy, Debug)]
3614pub struct StaticOrder2Atom<
3615    const N: usize,
3616    const H: usize,
3617    const GRADIENT_BITS: u128,
3618    const HESSIAN_BITS: u128,
3619> {
3620    value: f64,
3621    gradient: [f64; N],
3622    hessian: [f64; H],
3623}
3624
3625impl<const N: usize, const H: usize, const G: u128, const Q: u128> StaticOrder2Atom<N, H, G, Q> {
3626    /// Construct a generated atom. `H` must equal `N(N+1)/2`.
3627    #[inline(always)]
3628    #[must_use]
3629    pub fn new(value: f64, gradient: [f64; N], hessian: [f64; H]) -> Self {
3630        assert!(H == N * (N + 1) / 2, "invalid packed order-two shape");
3631        assert!(N <= 128 && H <= 128, "static atom sparsity mask overflow");
3632        Self {
3633            value,
3634            gradient,
3635            hessian,
3636        }
3637    }
3638
3639    /// Scalar value of the generated atom.
3640    #[inline(always)]
3641    #[must_use]
3642    pub fn value(&self) -> f64 {
3643        self.value
3644    }
3645
3646    /// Generated local gradient.
3647    #[inline(always)]
3648    #[must_use]
3649    pub fn gradient(&self) -> [f64; N] {
3650        self.gradient
3651    }
3652
3653    /// Generated local Hessian entry. The matrix is symmetric.
3654    #[inline(always)]
3655    #[must_use]
3656    pub fn hessian_at(&self, row: usize, column: usize) -> f64 {
3657        assert!(
3658            row < N && column < N,
3659            "static atom Hessian axis out of range"
3660        );
3661        let (row, column) = if row <= column {
3662            (row, column)
3663        } else {
3664            (column, row)
3665        };
3666        let index = row * (2 * N - row + 1) / 2 + column - row;
3667        self.hessian[index]
3668    }
3669}
3670
3671/// Order-two channel reader accepted by [`MappedOrder2Accumulator`].
3672///
3673/// Both ordinary forward jets and build-time-symbolic atoms implement this
3674/// interface. The accumulator owns the only global scatter/chain rule.
3675pub trait Order2AtomChannels<const N: usize> {
3676    /// Structurally live local gradient channels.
3677    const GRADIENT_BITS: u128;
3678    /// Structurally live packed upper-Hessian channels.
3679    const HESSIAN_BITS: u128;
3680    /// Local gradient entry.
3681    fn gradient_at(&self, axis: usize) -> f64;
3682    /// Local Hessian entry.
3683    fn hessian_at(&self, row: usize, column: usize) -> f64;
3684}
3685
3686impl<const N: usize> Order2AtomChannels<N> for Order2<N> {
3687    const GRADIENT_BITS: u128 = low_mask(N);
3688    const HESSIAN_BITS: u128 = low_mask(N * (N + 1) / 2);
3689
3690    #[inline(always)]
3691    fn gradient_at(&self, axis: usize) -> f64 {
3692        self.0.g[axis]
3693    }
3694
3695    #[inline(always)]
3696    fn hessian_at(&self, row: usize, column: usize) -> f64 {
3697        self.0.h[row][column]
3698    }
3699}
3700
3701impl<const N: usize, const H: usize, const G: u128, const Q: u128> Order2AtomChannels<N>
3702    for StaticOrder2Atom<N, H, G, Q>
3703{
3704    const GRADIENT_BITS: u128 = G;
3705    const HESSIAN_BITS: u128 = Q;
3706
3707    #[inline(always)]
3708    fn gradient_at(&self, axis: usize) -> f64 {
3709        self.gradient[axis]
3710    }
3711
3712    #[inline(always)]
3713    fn hessian_at(&self, row: usize, column: usize) -> f64 {
3714        StaticOrder2Atom::hessian_at(self, row, column)
3715    }
3716}
3717
3718const fn low_mask(channels: usize) -> u128 {
3719    if channels >= 128 {
3720        u128::MAX
3721    } else {
3722        (1u128 << channels) - 1
3723    }
3724}
3725
3726impl<const K: usize> MappedOrder2Accumulator<K> {
3727    /// Empty additive accumulator.
3728    #[inline(always)]
3729    #[must_use]
3730    pub fn zero() -> Self {
3731        Self {
3732            value: 0.0,
3733            gradient: [0.0; K],
3734            hessian: [[0.0; K]; K],
3735        }
3736    }
3737
3738    /// Scatter `f(atom)` using the exact order-two Faà di Bruno rule.
3739    ///
3740    /// `axes[i]` maps local derivative axis `i` to its global primary axis. The
3741    /// map must be injective. Repeated axes describe a non-injective linear
3742    /// pullback whose identified cross terms need multiplicity that this simple
3743    /// scatter deliberately does not represent.
3744    #[inline(always)]
3745    pub fn add_composed<const N: usize, const H: usize, A: Order2AtomChannels<N>>(
3746        &mut self,
3747        atom: &A,
3748        axes: [usize; N],
3749        derivatives: [f64; 3],
3750        value_add: bool,
3751        gradient_add: [bool; N],
3752        hessian_add: [bool; H],
3753    ) {
3754        assert!(H == N * (N + 1) / 2, "invalid mapped Hessian write shape");
3755        assert!(N <= 128 && H <= 128, "mapped atom sparsity mask overflow");
3756        assert!(
3757            axes.iter().all(|&axis| axis < K),
3758            "mapped atom axis must be within the global primary dimension"
3759        );
3760        assert!(
3761            axes.iter()
3762                .enumerate()
3763                .all(|(i, axis)| !axes[..i].contains(axis)),
3764            "mapped atom axes must be injective"
3765        );
3766
3767        if value_add {
3768            self.value += derivatives[0];
3769        } else {
3770            self.value = derivatives[0];
3771        }
3772        let mut packed = 0;
3773        for local_i in 0..N {
3774            let global_i = axes[local_i];
3775            if A::GRADIENT_BITS & (1u128 << local_i) != 0 {
3776                let channel = derivatives[1] * atom.gradient_at(local_i);
3777                if gradient_add[local_i] {
3778                    self.gradient[global_i] += channel;
3779                } else {
3780                    self.gradient[global_i] = channel;
3781                }
3782            }
3783            for local_j in local_i..N {
3784                let global_j = axes[local_j];
3785                let inner_live = A::HESSIAN_BITS & (1u128 << packed) != 0;
3786                let outer_live = A::GRADIENT_BITS & (1u128 << local_i) != 0
3787                    && A::GRADIENT_BITS & (1u128 << local_j) != 0;
3788                let channel = if inner_live {
3789                    let inner = derivatives[1] * atom.hessian_at(local_i, local_j);
3790                    if outer_live {
3791                        inner
3792                            + derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
3793                    } else {
3794                        inner
3795                    }
3796                } else if outer_live {
3797                    derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
3798                } else {
3799                    packed += 1;
3800                    continue;
3801                };
3802                if hessian_add[packed] {
3803                    self.hessian[global_i][global_j] += channel;
3804                    if global_i != global_j {
3805                        self.hessian[global_j][global_i] += channel;
3806                    }
3807                } else {
3808                    self.hessian[global_i][global_j] = channel;
3809                    if global_i != global_j {
3810                        self.hessian[global_j][global_i] = channel;
3811                    }
3812                }
3813                packed += 1;
3814            }
3815        }
3816    }
3817
3818    /// Finish the lowering in the standard row-kernel channel layout.
3819    #[inline(always)]
3820    #[must_use]
3821    pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
3822        (self.value, self.gradient, self.hessian)
3823    }
3824}
3825
3826/// One inner scalar in a runtime-width additive order-two composition.
3827///
3828/// Implementors expose only the inner gradient and Hessian. The corresponding
3829/// outer function's first and second derivatives live beside the source in the
3830/// caller's term type, so several terms that share an inner scalar may be fused
3831/// before entering [`DynamicOrder2Accumulator`].
3832pub trait DynamicOrder2Term {
3833    /// Outer first derivative `f'(q)` for this already-fused source.
3834    fn outer_first(&self) -> f64;
3835
3836    /// Outer second derivative `f''(q)` for this already-fused source.
3837    fn outer_second(&self) -> f64;
3838
3839    /// Inner first derivative `q_i`.
3840    fn inner_gradient(&self, axis: usize) -> f64;
3841
3842    /// Inner second derivative `q_ij`.
3843    fn inner_hessian(&self, row: usize, column: usize) -> f64;
3844}
3845
3846/// Allocation-minimal runtime-width lowering of an additive order-two program.
3847///
3848/// This is the dynamic-width sibling of [`MappedOrder2Accumulator`]. A caller
3849/// first reduces its mathematical row program to a scalar value plus `N`
3850/// independent composed sources. This accumulator then performs the universal
3851/// chain rule
3852///
3853/// ```text
3854/// g_i  = sum_t (f_t' q_t,i)
3855/// H_ij = sum_t (f_t'' q_t,i q_t,j + f_t' q_t,ij)
3856/// ```
3857///
3858/// in one gradient pass and one upper-triangular Hessian pass. It owns exactly
3859/// the two buffers that become the returned gradient and Hessian; no per-term
3860/// derivative buffer or jet temporary is allocated. `N` is const-generic so a
3861/// small fixed source set is visible to LLVM and can be unrolled.
3862#[derive(Debug)]
3863pub struct DynamicOrder2Accumulator {
3864    value: f64,
3865    gradient: Vec<f64>,
3866    hessian: Vec<f64>,
3867}
3868
3869impl DynamicOrder2Accumulator {
3870    /// Lower a fused additive source set at runtime primary width `dimension`.
3871    #[inline(always)]
3872    #[must_use]
3873    pub fn from_composed_sum<T: DynamicOrder2Term, const N: usize>(
3874        dimension: usize,
3875        value: f64,
3876        terms: &[T; N],
3877    ) -> Self {
3878        let mut gradient = vec![0.0; dimension];
3879        let mut hessian = vec![0.0; dimension * dimension];
3880
3881        for axis in 0..dimension {
3882            let mut channel = 0.0;
3883            for term in terms {
3884                channel += term.outer_first() * term.inner_gradient(axis);
3885            }
3886            gradient[axis] = channel;
3887        }
3888
3889        for row in 0..dimension {
3890            for column in row..dimension {
3891                let mut channel = 0.0;
3892                for term in terms {
3893                    let row_gradient = term.inner_gradient(row);
3894                    let column_gradient = term.inner_gradient(column);
3895                    channel += term.outer_second() * row_gradient * column_gradient
3896                        + term.outer_first() * term.inner_hessian(row, column);
3897                }
3898                hessian[row * dimension + column] = channel;
3899                hessian[column * dimension + row] = channel;
3900            }
3901        }
3902
3903        Self {
3904            value,
3905            gradient,
3906            hessian,
3907        }
3908    }
3909
3910    /// Finish the lowering as `(value, gradient, row-major symmetric Hessian)`.
3911    #[inline(always)]
3912    #[must_use]
3913    pub fn into_channels(self) -> (f64, Vec<f64>, Vec<f64>) {
3914        (self.value, self.gradient, self.hessian)
3915    }
3916}
3917
3918// ── Lane-batched Order-2 scalar: 4 rows per pass in SIMD lanes (perf) ────
3919//
3920// The hot per-row jet kernels evaluate ONE row's `(v, g, H)` tower at a time in
3921// scalar `f64`. A hand-written scalar derivative does the same. The throughput
3922// lever a jet has that scalar hand-code cannot is **row batching in SIMD
3923// lanes**: the order-≤2 Leibniz product `Order2::mul` is `O(K²)` independent
3924// per-channel float ops, and EVERY row runs the identical op graph on different
3925// data — the textbook SPMD shape. Packing `LANES = 4` rows into a `wide::f64x4`
3926// and running the algebra once per 4 rows replaces 4 scalar passes with one
3927// vector pass: the `K²` Hessian channel updates become `K²` NEON `.2d` / SSE2
3928// `pd` instructions covering 4 rows each, ~4× fewer FP instructions per row.
3929//
3930// The carried scalar field is abstracted by [`Lane`] so the SAME algebra body
3931// instantiates at `f64` (1 row, used as the bit-identity oracle) or
3932// [`wide::f64x4`] (4 rows). Bit-identity is structural, not approximate:
3933//
3934//   * Every arithmetic op is a plain lane-wise `+` / `-` / `*` (NEVER a fused
3935//     `mul_add`), and IEEE-754 double `+`/`-`/`*`/`/` are correctly rounded and
3936//     deterministic, so lane `i` of an `f64x4` op equals the scalar `f64` op on
3937//     that lane's inputs bit-for-bit.
3938//   * The transcendental derivative STACKS (`exp`/`ln`/`sqrt`/…) are produced
3939//     **per lane by the identical scalar code** ([`Lane::unary3`] unpacks, runs
3940//     the same `[f64; 3]` stack closure the scalar path runs, repacks), so the
3941//     only thing vectorised is the cheap rational tensor composition — the
3942//     library transcendental itself is the exact same `f64::exp` call per lane.
3943//   * The op order mirrors [`crate::jet_tower::Tower2`] term-for-term, so
3944//     [`Order2Lane<f64, K>`] is `to_bits`-identical to the production
3945//     [`Order2<K>`] (= `Tower2<K>`), and [`Order2Lane<f64x4, K>`] lane `i` is
3946//     `to_bits`-identical to that — proven by the `batch_tests` oracle below
3947//     (≥2000 random 4-row batches across `K ∈ {2,3,4,9}`).
3948
3949/// The scalar field a [`Order2Lane`] carries: either a single `f64` (one row,
3950/// the oracle) or a [`wide::f64x4`] (four rows evaluated in SIMD lanes). All ops
3951/// are plain lane-wise IEEE arithmetic, so a vector op equals the scalar op on
3952/// each lane bit-for-bit.
3953pub trait Lane: Copy {
3954    /// Broadcast a scalar to every lane.
3955    fn splat(x: f64) -> Self;
3956    /// Lane-wise `self + o`.
3957    fn add(self, o: Self) -> Self;
3958    /// Lane-wise `self - o`.
3959    fn sub(self, o: Self) -> Self;
3960    /// Lane-wise `self * o`.
3961    fn mul(self, o: Self) -> Self;
3962    /// The `f64` in lane `i` (`i < LANES`; `f64` ignores `i`).
3963    fn lane(self, i: usize) -> f64;
3964    /// Build the order-≤2 derivative stack `[f(u), f′(u), f″(u)]` **per lane**
3965    /// from the lane value `u`, via the SAME scalar `stack` closure the
3966    /// per-row path runs (so the transcendental/rational stack is bit-identical
3967    /// to the scalar evaluation — only the subsequent tensor composition is
3968    /// vectorised).
3969    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3];
3970    /// Build the order-≤4 derivative stack `[f, f′, f″, f‴, f⁗]` **per lane**
3971    /// from the lane value `u`, via the SAME scalar `stack` closure the per-row
3972    /// path runs. The one-/two-seed scalars ([`OneSeedLane`] / [`TwoSeedLane`])
3973    /// need outer derivatives one / two orders beyond their order-2 base, so
3974    /// they build their composition stack through this five-entry variant. As
3975    /// with [`unary3`](Lane::unary3), only the transcendental/rational stack is
3976    /// evaluated per lane (bit-identically to the scalar path); the subsequent
3977    /// tensor composition is vectorised.
3978    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5];
3979}
3980
3981impl Lane for f64 {
3982    #[inline]
3983    fn splat(x: f64) -> Self {
3984        x
3985    }
3986    #[inline]
3987    fn add(self, o: Self) -> Self {
3988        self + o
3989    }
3990    #[inline]
3991    fn sub(self, o: Self) -> Self {
3992        self - o
3993    }
3994    #[inline]
3995    fn mul(self, o: Self) -> Self {
3996        self * o
3997    }
3998    #[inline]
3999    fn lane(self, _: usize) -> f64 {
4000        self
4001    }
4002    #[inline]
4003    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
4004        stack(self)
4005    }
4006    #[inline]
4007    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
4008        stack(self)
4009    }
4010}
4011
4012impl Lane for wide::f64x4 {
4013    #[inline]
4014    fn splat(x: f64) -> Self {
4015        wide::f64x4::splat(x)
4016    }
4017    #[inline]
4018    fn add(self, o: Self) -> Self {
4019        self + o
4020    }
4021    #[inline]
4022    fn sub(self, o: Self) -> Self {
4023        self - o
4024    }
4025    #[inline]
4026    fn mul(self, o: Self) -> Self {
4027        self * o
4028    }
4029    #[inline]
4030    fn lane(self, i: usize) -> f64 {
4031        self.to_array()[i]
4032    }
4033    #[inline]
4034    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
4035        let a = self.to_array();
4036        let mut d0 = [0.0_f64; 4];
4037        let mut d1 = [0.0_f64; 4];
4038        let mut d2 = [0.0_f64; 4];
4039        for i in 0..4 {
4040            let s = stack(a[i]);
4041            d0[i] = s[0];
4042            d1[i] = s[1];
4043            d2[i] = s[2];
4044        }
4045        [
4046            wide::f64x4::new(d0),
4047            wide::f64x4::new(d1),
4048            wide::f64x4::new(d2),
4049        ]
4050    }
4051    #[inline]
4052    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
4053        let a = self.to_array();
4054        let mut d = [[0.0_f64; 4]; 5];
4055        for i in 0..4 {
4056            let s = stack(a[i]);
4057            for (k, dk) in d.iter_mut().enumerate() {
4058                dk[i] = s[k];
4059            }
4060        }
4061        [
4062            wide::f64x4::new(d[0]),
4063            wide::f64x4::new(d[1]),
4064            wide::f64x4::new(d[2]),
4065            wide::f64x4::new(d[3]),
4066            wide::f64x4::new(d[4]),
4067        ]
4068    }
4069}
4070
4071/// A lane-batched order-≤2 Taylor scalar: value / gradient / Hessian carried in
4072/// a SIMD field [`L: Lane`](Lane). With `L = f64x4` one instance carries FOUR
4073/// rows at once, so the row loop processes 4 rows per vector pass instead of one
4074/// per scalar pass.
4075///
4076/// The channel layout and every float op mirror [`crate::jet_tower::Tower2`]
4077/// term-for-term, so `Order2Lane<f64, K>` is `to_bits`-identical to the
4078/// production [`Order2<K>`] and `Order2Lane<f64x4, K>` lane `i` is
4079/// `to_bits`-identical to that (see the module note and `batch_tests`).
4080#[derive(Clone, Copy, Debug)]
4081pub struct Order2Lane<L: Lane, const K: usize> {
4082    /// Value channel `ℓ` (one entry per lane/row).
4083    pub v: L,
4084    /// Gradient channel `∂ℓ/∂p_a`.
4085    pub g: [L; K],
4086    /// Hessian channel `∂²ℓ/∂p_a∂p_b` (symmetric).
4087    pub h: [[L; K]; K],
4088}
4089
4090/// The 4-rows-per-pass batched order-≤2 scalar (`wide::f64x4` lanes).
4091pub type Order2Batch<const K: usize> = Order2Lane<wide::f64x4, K>;
4092
4093impl<L: Lane, const K: usize> Order2Lane<L, K> {
4094    /// A constant: value `c` in every channel-zero slot.
4095    #[inline]
4096    pub fn constant(c: L) -> Self {
4097        Order2Lane {
4098            v: c,
4099            g: [L::splat(0.0); K],
4100            h: [[L::splat(0.0); K]; K],
4101        }
4102    }
4103
4104    /// The seeded variable `p_axis` at (per-lane) value `value`: unit first
4105    /// derivative in slot `axis`. With `L = f64x4`, `value` packs the four
4106    /// rows' values of primary `axis`.
4107    #[inline]
4108    pub fn variable(value: L, axis: usize) -> Self {
4109        let mut out = Self::constant(value);
4110        out.g[axis] = L::splat(1.0);
4111        out
4112    }
4113
4114    /// Lane-wise `self + o` (mirrors `Tower2` Add: per-channel add).
4115    #[inline]
4116    pub fn add(&self, o: &Self) -> Self {
4117        let mut out = *self;
4118        out.v = self.v.add(o.v);
4119        for i in 0..K {
4120            out.g[i] = self.g[i].add(o.g[i]);
4121            for j in 0..K {
4122                out.h[i][j] = self.h[i][j].add(o.h[i][j]);
4123            }
4124        }
4125        out
4126    }
4127
4128    /// Multiply every channel by the plain scalar `s` (mirrors `Tower2::scale`).
4129    #[inline]
4130    pub fn scale(&self, s: f64) -> Self {
4131        let sl = L::splat(s);
4132        let mut out = *self;
4133        out.v = self.v.mul(sl);
4134        for i in 0..K {
4135            out.g[i] = self.g[i].mul(sl);
4136            for j in 0..K {
4137                out.h[i][j] = self.h[i][j].mul(sl);
4138            }
4139        }
4140        out
4141    }
4142
4143    /// Lane-wise `self - o`, expressed as `self + o·(-1)` exactly as
4144    /// [`Order2::sub`] / `Tower4::sub` do, so signed-zero handling matches.
4145    #[inline]
4146    pub fn sub(&self, o: &Self) -> Self {
4147        self.add(&o.scale(-1.0))
4148    }
4149
4150    /// Negate every channel (= `scale(-1.0)`, matching [`Order2::neg`]).
4151    #[inline]
4152    pub fn neg(&self) -> Self {
4153        self.scale(-1.0)
4154    }
4155
4156    /// Exact order-≤2 Leibniz product, term-for-term identical to
4157    /// [`crate::jet_tower::Tower2::mul`] (same factor order, no `mul_add`).
4158    ///
4159    /// The Hessian channel is symmetric under `i ↔ j` (see
4160    /// [`crate::jet_tower::Tower2::mul`] for why the invariant always holds), so
4161    /// we compute the upper triangle `j ≥ i` and mirror it — `K(K+1)/2` lane
4162    /// entry-chains instead of `K²`. Because each lane entry is already a full
4163    /// SIMD op (no cross-`j` lane packing to lose), halving the entry count is a
4164    /// direct throughput win (~18 % on `Order2Lane<f64x4, 9>`, the survival batch
4165    /// kernel, and ~2× on the `f64` oracle). The upper triangle uses the EXACT
4166    /// term order of `Tower2::mul`, so `Order2Lane<f64>` stays `to_bits`-identical
4167    /// to `Order2` (= `Tower2`) and `Order2Lane<f64x4>` lane `i` stays
4168    /// `to_bits`-identical to that; the mirror makes the batch Hessian exactly
4169    /// symmetric, matching the scalar `Tower2::mul` (which mirrors identically).
4170    #[inline]
4171    pub fn mul(&self, o: &Self) -> Self {
4172        let a = self;
4173        let b = o;
4174        let mut out = Self::constant(a.v.mul(b.v));
4175        for i in 0..K {
4176            // a.v*b.g[i] + a.g[i]*b.v
4177            out.g[i] = a.v.mul(b.g[i]).add(a.g[i].mul(b.v));
4178        }
4179        for i in 0..K {
4180            for j in i..K {
4181                // a.v*b.h + a.g[i]*b.g[j] + a.g[j]*b.g[i] + a.h*b.v
4182                let hij =
4183                    a.v.mul(b.h[i][j])
4184                        .add(a.g[i].mul(b.g[j]))
4185                        .add(a.g[j].mul(b.g[i]))
4186                        .add(a.h[i][j].mul(b.v));
4187                out.h[i][j] = hij;
4188                out.h[j][i] = hij;
4189            }
4190        }
4191        out
4192    }
4193
4194    /// Exact order-≤2 Faà di Bruno composition `f ∘ self`, given the per-lane
4195    /// derivative stack `d = [f(u), f′(u), f″(u)]`. Mirrors
4196    /// [`crate::jet_tower::Tower2::compose_unary`] term-for-term (`acc` starts at
4197    /// `0` then accumulates, so signed-zero collapses identically).
4198    #[inline]
4199    pub fn compose_unary(&self, d: [L; 3]) -> Self {
4200        let mut out = Self::constant(d[0]);
4201        for i in 0..K {
4202            let mut acc = L::splat(0.0);
4203            acc = acc.add(d[1].mul(self.g[i]));
4204            out.g[i] = acc;
4205        }
4206        for i in 0..K {
4207            for j in 0..K {
4208                let mut acc = L::splat(0.0);
4209                acc = acc.add(d[1].mul(self.h[i][j]));
4210                acc = acc.add(d[2].mul(self.g[i]).mul(self.g[j]));
4211                out.h[i][j] = acc;
4212            }
4213        }
4214        out
4215    }
4216
4217    /// `e^self`, per-lane stack `[e, e, e]` (matches the [`JetScalar::exp`]
4218    /// default forwarded through `Order2`).
4219    #[inline]
4220    pub fn exp(&self) -> Self {
4221        let d = self.v.unary3(|u| {
4222            let e = u.exp();
4223            [e, e, e]
4224        });
4225        self.compose_unary(d)
4226    }
4227
4228    /// `ln(self)`; caller guarantees positivity. Per-lane stack
4229    /// `[ln u, 1/u, -1/u²]` (matches [`JetScalar::ln`] truncated to order 2).
4230    #[inline]
4231    pub fn ln(&self) -> Self {
4232        let d = self.v.unary3(|u| {
4233            let r = 1.0 / u;
4234            [u.ln(), r, -r * r]
4235        });
4236        self.compose_unary(d)
4237    }
4238
4239    /// `√self`; caller guarantees positivity. Per-lane stack
4240    /// `[s, 0.5/s, -0.25/(u·s)]` (matches [`JetScalar::sqrt`]).
4241    #[inline]
4242    pub fn sqrt(&self) -> Self {
4243        let d = self.v.unary3(|u| {
4244            let s = u.sqrt();
4245            [s, 0.5 / s, -0.25 / (u * s)]
4246        });
4247        self.compose_unary(d)
4248    }
4249
4250    /// `1/self`. Per-lane stack `[r, -r², 2r³]` (matches [`JetScalar::recip`]).
4251    #[inline]
4252    pub fn recip(&self) -> Self {
4253        let d = self.v.unary3(|u| {
4254            let r = 1.0 / u;
4255            let r2 = r * r;
4256            [r, -r2, 2.0 * r2 * r]
4257        });
4258        self.compose_unary(d)
4259    }
4260
4261    /// `self^a` for real `a`; caller guarantees a positive base. Per-lane
4262    /// falling-factorial stack (matches [`JetScalar::powf`]).
4263    #[inline]
4264    pub fn powf(&self, a: f64) -> Self {
4265        let d = self.v.unary3(|u| {
4266            [
4267                u.powf(a),
4268                a * u.powf(a - 1.0),
4269                a * (a - 1.0) * u.powf(a - 2.0),
4270            ]
4271        });
4272        self.compose_unary(d)
4273    }
4274}
4275
4276impl<const K: usize> Order2Batch<K> {
4277    /// Extract lane `i`'s `(v, g, H)` as a production [`Order2<K>`] scalar.
4278    /// Lane `i` is `to_bits`-identical to evaluating the same program at
4279    /// [`Order2<K>`] on row `i` (see `batch_tests`).
4280    #[inline]
4281    #[must_use]
4282    pub fn lane(&self, i: usize) -> Order2<K> {
4283        let mut t = crate::jet_tower::Tower2::<K>::constant(self.v.lane(i));
4284        for a in 0..K {
4285            t.g[a] = self.g[a].lane(i);
4286            for b in 0..K {
4287                t.h[a][b] = self.h[a][b].lane(i);
4288            }
4289        }
4290        Order2(t)
4291    }
4292}
4293
4294// ── Order1<K>: value / gradient only (doc §A.1, first-order prune) ──────
4295
4296/// Truncated FIRST-order scalar: value `v` and gradient `g_a` only — NO Hessian.
4297///
4298/// This is [`Order2`] with the K×K Hessian channel deleted. Its value and
4299/// gradient are computed by the SAME order-≤1 truncation of the Leibniz / Faà
4300/// di Bruno rules that [`Order2`] uses for those two channels, with the float
4301/// operations applied in the identical order — so its `(v, g)` is BIT-IDENTICAL
4302/// to both [`Order2`]'s and a full [`crate::jet_tower::Tower4`]'s order-≤1
4303/// channels. Use it at a consumer that reads ONLY value + gradient (the SAE
4304/// β-border channel: the reconstruction is linear in β, so the Hessian-in-β
4305/// vanishes and the dense K×K Hessian product `Tower2::mul` would build is pure
4306/// discarded work). Order-≤1 value/gradient never read any input's Hessian, so
4307/// dropping that channel changes neither result nor float-op order — it only
4308/// removes the `K²` arithmetic that produced an unread tensor.
4309#[derive(Clone, Copy, Debug)]
4310pub struct Order1<const K: usize> {
4311    /// Value ℓ.
4312    pub v: f64,
4313    /// Gradient ∂ℓ/∂p_a.
4314    pub g: [f64; K],
4315}
4316
4317impl<const K: usize> Order1<K> {
4318    /// Read the gradient channel `g_a = ∂ℓ/∂p_a`.
4319    #[inline]
4320    #[must_use]
4321    pub fn g(&self) -> &[f64; K] {
4322        &self.g
4323    }
4324
4325    /// Consume the jet and move out its value and gradient channels.
4326    #[inline]
4327    #[must_use]
4328    pub fn into_channels(self) -> (f64, [f64; K]) {
4329        (self.v, self.g)
4330    }
4331}
4332
4333impl<const K: usize> JetScalar<K> for Order1<K> {
4334    fn constant(c: f64) -> Self {
4335        // Order2::constant -> Tower2::constant: value c, all derivatives zero.
4336        Order1 { v: c, g: [0.0; K] }
4337    }
4338    fn variable(x: f64, axis: usize) -> Self {
4339        // Order2::variable -> Tower2::variable: unit first derivative in `axis`.
4340        let mut g = [0.0; K];
4341        g[axis] = 1.0;
4342        Order1 { v: x, g }
4343    }
4344}
4345
4346impl<const K: usize> crate::nested_dual::JetField for Order1<K> {
4347    fn value(&self) -> f64 {
4348        self.v
4349    }
4350    fn add(&self, o: &Self) -> Self {
4351        // Tower2 Add: out.v += o.v; out.g[i] += o.g[i] (same float order).
4352        let mut g = self.g;
4353        for i in 0..K {
4354            g[i] += o.g[i];
4355        }
4356        Order1 { v: self.v + o.v, g }
4357    }
4358    fn sub(&self, o: &Self) -> Self {
4359        // Mirror Order2::sub == self + o.scale(-1.0) exactly: scale then add.
4360        self.add(&o.scale(-1.0))
4361    }
4362    fn mul(&self, o: &Self) -> Self {
4363        // Tower2::mul value/grad terms, identical float order:
4364        //   v = a.v*b.v;  g[i] = a.v*b.g[i] + a.g[i]*b.v.
4365        // (The Hessian loop `a.v*b.h + a.g*b.g + ... + a.h*b.v` is the discarded
4366        //  work this type exists to skip; it never feeds v or g.)
4367        let a = self;
4368        let b = o;
4369        let mut g = [0.0; K];
4370        for i in 0..K {
4371            g[i] = a.v * b.g[i] + a.g[i] * b.v;
4372        }
4373        Order1 { v: a.v * b.v, g }
4374    }
4375    fn neg(&self) -> Self {
4376        // Order2::neg == self.0.scale(-1.0).
4377        self.scale(-1.0)
4378    }
4379    fn scale(&self, s: f64) -> Self {
4380        // Tower2::scale: out.v *= s; out.g[i] *= s (same float order).
4381        let mut g = self.g;
4382        for i in 0..K {
4383            g[i] *= s;
4384        }
4385        Order1 { v: self.v * s, g }
4386    }
4387    fn compose_unary(&self, d: [f64; 5]) -> Self {
4388        // Faà di Bruno truncated to order ≤ 1 (matches `faa_di_bruno` /
4389        // `Tower2::compose_unary` for the value and gradient channels):
4390        //   value channel (m=0): d[0].
4391        //   grad channel (positions=[i], single partition {{0}}): d[1]·g[i].
4392        // Order-≤1 reads only d[0], d[1]; trailing stack entries are unused.
4393        let mut g = [0.0; K];
4394        for i in 0..K {
4395            g[i] = d[1] * self.g[i];
4396        }
4397        Order1 { v: d[0], g }
4398    }
4399}
4400
4401// ── OneSeed<K>: one-seed directional, contracted third (doc §A.2) ───────
4402
4403/// One-seed directional scalar: an [`Order2`] base plus ONE nilpotent ε
4404/// (`ε² = 0`) whose coefficient is itself an [`Order2`].
4405///
4406/// A scalar is `s = base + ε·eps`. Arithmetic is the `ε² = 0` truncation of the
4407/// product (doc §A.2): the base parts multiply as ordinary [`Order2`] products,
4408/// and the ε-coefficient picks up `a.base·b.eps + a.eps·b.base`. Composition
4409/// pushes ε through one extra outer derivative.
4410///
4411/// Seed each primary with [`seed_direction`](Self::seed_direction): the base is
4412/// the usual seeded variable (carrying `e_a` for the Hessian channel) and the
4413/// ε-coefficient is the FIXED contraction direction `u_a` (a constant). Then the
4414/// ε-component of the evaluated Hessian channel is the contracted third
4415/// `[eps.h][a][b] = Σ_c ℓ_{abc} u_c` — exactly `row_third_contracted(dir = u)`,
4416/// without materialising `t3`.
4417#[derive(Clone, Copy, Debug)]
4418pub struct OneSeed<const K: usize> {
4419    /// The `ε⁰` part: value / gradient / Hessian of `ℓ`.
4420    pub base: Order2<K>,
4421    /// The `ε¹` part: value / gradient / Hessian of the ε-coefficient. After a
4422    /// `seed_direction(u)` evaluation, `eps.h[a][b] = Σ_c ℓ_{abc} u_c`.
4423    pub eps: Order2<K>,
4424}
4425
4426impl<const K: usize> OneSeed<K> {
4427    /// Seed primary `axis` at value `x` with ε-direction component `u_axis`:
4428    /// `p_axis = p_axis⁰ + x-seed + ε·u_axis`, i.e. base = `variable(x, axis)`
4429    /// and eps = `constant(u_axis)` (doc §A.2 "Seeding").
4430    pub fn seed_direction(x: f64, axis: usize, u_axis: f64) -> Self {
4431        OneSeed {
4432            base: Order2::variable(x, axis),
4433            eps: Order2::constant(u_axis),
4434        }
4435    }
4436
4437    /// The contracted-third channel after a `seed_direction(u)` evaluation:
4438    /// `out[a][b] = Σ_c ℓ_{abc} u_c`, i.e. the ε-coefficient's Hessian (doc §A.2).
4439    pub fn contracted_third(&self) -> [[f64; K]; K] {
4440        *self.eps.h()
4441    }
4442}
4443
4444impl<const K: usize> JetScalar<K> for OneSeed<K> {
4445    fn constant(c: f64) -> Self {
4446        OneSeed {
4447            base: Order2::constant(c),
4448            eps: Order2::constant(0.0),
4449        }
4450    }
4451    fn variable(x: f64, axis: usize) -> Self {
4452        // No ε-direction unless seeded via `seed_direction`.
4453        OneSeed {
4454            base: Order2::variable(x, axis),
4455            eps: Order2::constant(0.0),
4456        }
4457    }
4458}
4459
4460impl<const K: usize> crate::nested_dual::JetField for OneSeed<K> {
4461    fn value(&self) -> f64 {
4462        self.base.value()
4463    }
4464    fn add(&self, o: &Self) -> Self {
4465        OneSeed {
4466            base: self.base.add(&o.base),
4467            eps: self.eps.add(&o.eps),
4468        }
4469    }
4470    fn sub(&self, o: &Self) -> Self {
4471        OneSeed {
4472            base: self.base.sub(&o.base),
4473            eps: self.eps.sub(&o.eps),
4474        }
4475    }
4476    fn mul(&self, o: &Self) -> Self {
4477        // (a.base + ε a.eps)(b.base + ε b.eps), dropping ε².  Build
4478        // the epsilon channels directly instead of materialising two complete
4479        // Order2 products and adding them.  The upper triangle is sufficient:
4480        // every input Hessian is symmetric and this polynomial preserves that
4481        // invariant exactly.
4482        let ab = &self.base.0;
4483        let ae = &self.eps.0;
4484        let bb = &o.base.0;
4485        let be = &o.eps.0;
4486        let mut eps = crate::jet_tower::Tower2::<K>::zero();
4487        eps.v = ab.v * be.v + ae.v * bb.v;
4488        for i in 0..K {
4489            eps.g[i] = ab.v * be.g[i] + ab.g[i] * be.v + ae.v * bb.g[i] + ae.g[i] * bb.v;
4490        }
4491        for i in 0..K {
4492            for j in i..K {
4493                let channel = ab.v * be.h[i][j]
4494                    + ab.g[i] * be.g[j]
4495                    + ab.g[j] * be.g[i]
4496                    + ab.h[i][j] * be.v
4497                    + ae.v * bb.h[i][j]
4498                    + ae.g[i] * bb.g[j]
4499                    + ae.g[j] * bb.g[i]
4500                    + ae.h[i][j] * bb.v;
4501                eps.h[i][j] = channel;
4502                eps.h[j][i] = channel;
4503            }
4504        }
4505        OneSeed {
4506            base: self.base.mul(&o.base),
4507            eps: Order2(eps),
4508        }
4509    }
4510    fn neg(&self) -> Self {
4511        OneSeed {
4512            base: self.base.neg(),
4513            eps: self.eps.neg(),
4514        }
4515    }
4516    fn scale(&self, s: f64) -> Self {
4517        OneSeed {
4518            base: self.base.scale(s),
4519            eps: self.eps.scale(s),
4520        }
4521    }
4522    fn compose_unary(&self, d: [f64; 5]) -> Self {
4523        // f(base + ε eps) = f(base) + ε · f'(base)·eps (ε² = 0).
4524        // Fuse the shifted composition and product into their final channels;
4525        // this removes one temporary Order2 composition and one Order2 product
4526        // from every unary primitive in a directional row program.
4527        let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
4528        let b = &self.base.0;
4529        let e = &self.eps.0;
4530        let mut eps = crate::jet_tower::Tower2::<K>::zero();
4531        eps.v = d[1] * e.v;
4532        for i in 0..K {
4533            eps.g[i] = d[2] * b.g[i] * e.v + d[1] * e.g[i];
4534        }
4535        for i in 0..K {
4536            for j in i..K {
4537                let channel = d[1] * e.h[i][j]
4538                    + d[2] * (b.g[i] * e.g[j] + b.g[j] * e.g[i] + b.h[i][j] * e.v)
4539                    + d[3] * b.g[i] * b.g[j] * e.v;
4540                eps.h[i][j] = channel;
4541                eps.h[j][i] = channel;
4542            }
4543        }
4544        OneSeed {
4545            base,
4546            eps: Order2(eps),
4547        }
4548    }
4549}
4550
4551// ── OneSeedLane<L, K>: lane-batched one-seed directional (doc §A.2) ──────
4552
4553/// Lane-batched [`OneSeed`]: the same one-seed directional scalar with its two
4554/// [`Order2`] parts re-typed to [`Order2Lane<L, K>`], so one `L = f64x4`
4555/// instance carries FOUR rows' contracted-third evaluations per vector pass.
4556///
4557/// Every operation (`add`/`sub`/`mul`/`neg`/`scale`/`compose_unary` and the
4558/// transcendentals) is a term-for-term structural re-type of the scalar
4559/// [`OneSeed`] ops onto the lane-implemented [`Order2Lane`] algebra. With
4560/// `L = f64`, `OneSeedLane<f64, K>` is `to_bits`-identical to [`OneSeed<K>`];
4561/// with `L = f64x4`, lane `i` is `to_bits`-identical to that (see `batch_tests`).
4562#[derive(Clone, Copy, Debug)]
4563pub struct OneSeedLane<L: Lane, const K: usize> {
4564    /// The `ε⁰` part (lane-batched value / gradient / Hessian of `ℓ`).
4565    pub base: Order2Lane<L, K>,
4566    /// The `ε¹` part. After a `seed_direction(u)` evaluation,
4567    /// `eps.h[a][b]` lane `i` is row `i`'s `Σ_c ℓ_{abc} u_c`.
4568    pub eps: Order2Lane<L, K>,
4569}
4570
4571/// The 4-rows-per-pass batched one-seed scalar (`wide::f64x4` lanes).
4572pub type OneSeedBatch<const K: usize> = OneSeedLane<wide::f64x4, K>;
4573
4574impl<L: Lane, const K: usize> OneSeedLane<L, K> {
4575    /// A constant: base = `constant(c)`, ε-part zero (mirrors [`OneSeed::constant`]).
4576    #[inline]
4577    pub fn constant(c: L) -> Self {
4578        OneSeedLane {
4579            base: Order2Lane::constant(c),
4580            eps: Order2Lane::constant(L::splat(0.0)),
4581        }
4582    }
4583
4584    /// The seeded variable `p_axis` at (per-lane) value `value`, no ε-direction
4585    /// (mirrors [`OneSeed::variable`]).
4586    #[inline]
4587    pub fn variable(value: L, axis: usize) -> Self {
4588        OneSeedLane {
4589            base: Order2Lane::variable(value, axis),
4590            eps: Order2Lane::constant(L::splat(0.0)),
4591        }
4592    }
4593
4594    /// Seed primary `axis` at (per-lane) value `value` with ε-direction
4595    /// component `u_axis`: base = `variable(value, axis)`, eps = `constant(u_axis)`
4596    /// (mirrors [`OneSeed::seed_direction`]). With `L = f64x4`, `value` / `u_axis`
4597    /// pack the four rows' values / directions of primary `axis`.
4598    #[inline]
4599    pub fn seed_direction(value: L, axis: usize, u_axis: L) -> Self {
4600        OneSeedLane {
4601            base: Order2Lane::variable(value, axis),
4602            eps: Order2Lane::constant(u_axis),
4603        }
4604    }
4605
4606    /// The contracted-third channel after a `seed_direction(u)` evaluation:
4607    /// `out[a][b]` lane `i` is row `i`'s `Σ_c ℓ_{abc} u_c` (the ε-part Hessian).
4608    #[inline]
4609    #[must_use]
4610    pub fn contracted_third(&self) -> [[L; K]; K] {
4611        self.eps.h
4612    }
4613
4614    /// Lane-wise `self + o` (mirrors [`OneSeed::add`]).
4615    #[inline]
4616    pub fn add(&self, o: &Self) -> Self {
4617        OneSeedLane {
4618            base: self.base.add(&o.base),
4619            eps: self.eps.add(&o.eps),
4620        }
4621    }
4622
4623    /// Lane-wise `self - o` (mirrors [`OneSeed::sub`]).
4624    #[inline]
4625    pub fn sub(&self, o: &Self) -> Self {
4626        OneSeedLane {
4627            base: self.base.sub(&o.base),
4628            eps: self.eps.sub(&o.eps),
4629        }
4630    }
4631
4632    /// Lane-wise `self · o`, ε² = 0 truncation (mirrors [`OneSeed::mul`]).
4633    #[inline]
4634    pub fn mul(&self, o: &Self) -> Self {
4635        let ab = &self.base;
4636        let ae = &self.eps;
4637        let bb = &o.base;
4638        let be = &o.eps;
4639        let mut eps = Order2Lane::constant(ab.v.mul(be.v).add(ae.v.mul(bb.v)));
4640        for i in 0..K {
4641            eps.g[i] =
4642                ab.v.mul(be.g[i])
4643                    .add(ab.g[i].mul(be.v))
4644                    .add(ae.v.mul(bb.g[i]))
4645                    .add(ae.g[i].mul(bb.v));
4646        }
4647        for i in 0..K {
4648            for j in i..K {
4649                let channel =
4650                    ab.v.mul(be.h[i][j])
4651                        .add(ab.g[i].mul(be.g[j]))
4652                        .add(ab.g[j].mul(be.g[i]))
4653                        .add(ab.h[i][j].mul(be.v))
4654                        .add(ae.v.mul(bb.h[i][j]))
4655                        .add(ae.g[i].mul(bb.g[j]))
4656                        .add(ae.g[j].mul(bb.g[i]))
4657                        .add(ae.h[i][j].mul(bb.v));
4658                eps.h[i][j] = channel;
4659                eps.h[j][i] = channel;
4660            }
4661        }
4662        OneSeedLane {
4663            base: self.base.mul(&o.base),
4664            eps,
4665        }
4666    }
4667
4668    /// Negate every part (mirrors [`OneSeed::neg`]).
4669    #[inline]
4670    pub fn neg(&self) -> Self {
4671        OneSeedLane {
4672            base: self.base.neg(),
4673            eps: self.eps.neg(),
4674        }
4675    }
4676
4677    /// Multiply every part by the plain scalar `s` (mirrors [`OneSeed::scale`]).
4678    #[inline]
4679    pub fn scale(&self, s: f64) -> Self {
4680        OneSeedLane {
4681            base: self.base.scale(s),
4682            eps: self.eps.scale(s),
4683        }
4684    }
4685
4686    /// Exact order-≤2-per-part Faà di Bruno composition `f ∘ self`, given the
4687    /// per-lane outer-derivative stack `d = [f, f′, f″, f‴, f⁗]`. Term-for-term
4688    /// identical to the fused [`OneSeed::compose_unary`] channels: the base reads
4689    /// `d[0..=2]`, while the ε-coefficient directly evaluates
4690    /// `f′(base)·eps` from `d[1..=3]` without an intermediate jet.
4691    #[inline]
4692    pub fn compose_unary(&self, d: [L; 5]) -> Self {
4693        let base = self.base.compose_unary([d[0], d[1], d[2]]);
4694        let b = &self.base;
4695        let e = &self.eps;
4696        let mut eps = Order2Lane::constant(d[1].mul(e.v));
4697        for i in 0..K {
4698            eps.g[i] = d[2].mul(b.g[i]).mul(e.v).add(d[1].mul(e.g[i]));
4699        }
4700        for i in 0..K {
4701            for j in i..K {
4702                let mixed = b.g[i]
4703                    .mul(e.g[j])
4704                    .add(b.g[j].mul(e.g[i]))
4705                    .add(b.h[i][j].mul(e.v));
4706                let channel = d[1]
4707                    .mul(e.h[i][j])
4708                    .add(d[2].mul(mixed))
4709                    .add(d[3].mul(b.g[i]).mul(b.g[j]).mul(e.v));
4710                eps.h[i][j] = channel;
4711                eps.h[j][i] = channel;
4712            }
4713        }
4714        OneSeedLane { base, eps }
4715    }
4716
4717    /// `e^self`, per-lane stack `[e, e, e, e, e]` (matches [`JetScalar::exp`]).
4718    #[inline]
4719    pub fn exp(&self) -> Self {
4720        let d = self.base.v.unary5(|u| {
4721            let e = u.exp();
4722            [e, e, e, e, e]
4723        });
4724        self.compose_unary(d)
4725    }
4726
4727    /// `ln(self)`; caller guarantees positivity (matches [`JetScalar::ln`]).
4728    #[inline]
4729    pub fn ln(&self) -> Self {
4730        let d = self.base.v.unary5(|u| {
4731            let r = 1.0 / u;
4732            [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
4733        });
4734        self.compose_unary(d)
4735    }
4736
4737    /// `√self`; caller guarantees positivity (matches [`JetScalar::sqrt`]).
4738    #[inline]
4739    pub fn sqrt(&self) -> Self {
4740        let d = self.base.v.unary5(|u| {
4741            let s = u.sqrt();
4742            [
4743                s,
4744                0.5 / s,
4745                -0.25 / (u * s),
4746                0.375 / (u * u * s),
4747                -0.9375 / (u * u * u * s),
4748            ]
4749        });
4750        self.compose_unary(d)
4751    }
4752
4753    /// `1/self` (matches [`JetScalar::recip`]).
4754    #[inline]
4755    pub fn recip(&self) -> Self {
4756        let d = self.base.v.unary5(|u| {
4757            let r = 1.0 / u;
4758            let r2 = r * r;
4759            [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
4760        });
4761        self.compose_unary(d)
4762    }
4763
4764    /// `self^a` for real `a`; caller guarantees a positive base (matches
4765    /// [`JetScalar::powf`]).
4766    #[inline]
4767    pub fn powf(&self, a: f64) -> Self {
4768        let d = self.base.v.unary5(|u| {
4769            [
4770                u.powf(a),
4771                a * u.powf(a - 1.0),
4772                a * (a - 1.0) * u.powf(a - 2.0),
4773                a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
4774                a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
4775            ]
4776        });
4777        self.compose_unary(d)
4778    }
4779
4780    /// `ln Γ(self)`; caller guarantees positivity (matches [`JetScalar::ln_gamma`],
4781    /// same hand-certified stack).
4782    #[inline]
4783    pub fn ln_gamma(&self) -> Self {
4784        let d = self
4785            .base
4786            .v
4787            .unary5(crate::jet_tower::ln_gamma_derivative_stack);
4788        self.compose_unary(d)
4789    }
4790
4791    /// `ψ(self)` digamma; caller guarantees positivity (matches
4792    /// [`JetScalar::digamma`], same hand-certified stack).
4793    #[inline]
4794    pub fn digamma(&self) -> Self {
4795        let d = self
4796            .base
4797            .v
4798            .unary5(crate::jet_tower::digamma_derivative_stack);
4799        self.compose_unary(d)
4800    }
4801}
4802
4803impl<const K: usize> OneSeedBatch<K> {
4804    /// Extract lane `i`'s parts as a production [`OneSeed<K>`]. Lane `i` is
4805    /// `to_bits`-identical to evaluating the same program at [`OneSeed<K>`] on
4806    /// row `i` (see `batch_tests`).
4807    #[inline]
4808    #[must_use]
4809    pub fn lane(&self, i: usize) -> OneSeed<K> {
4810        OneSeed {
4811            base: self.base.lane(i),
4812            eps: self.eps.lane(i),
4813        }
4814    }
4815}
4816
4817// ── TwoSeed<K>: two-seed, contracted fourth (doc §A.3) ──────────────────
4818
4819/// Two-seed scalar: an [`Order2`] base plus TWO nilpotents ε, δ
4820/// (`ε² = δ² = 0`, `εδ` retained) — four [`Order2`] parts
4821/// `s = base + ε·eps + δ·del + εδ·eps_del`.
4822///
4823/// Product truncates `ε² = δ² = 0` (doc §A.3): each part is built from
4824/// [`Order2`] products of the four input parts. Composition picks up
4825/// successively higher outer derivatives, the cross part carrying the second
4826/// Faà di Bruno term `f''·eps·del + f'·eps_del`.
4827///
4828/// Seed each primary with [`seed`](Self::seed): base = `variable(x, axis)`,
4829/// eps = `constant(u_axis)`, del = `constant(v_axis)`, eps_del = `constant(0)`.
4830/// Then the εδ-component of the evaluated Hessian channel is the contracted
4831/// fourth `[eps_del.h][a][b] = Σ_{cd} ℓ_{abcd} u_c v_d` — exactly
4832/// `row_fourth_contracted(u, v)`, without materialising `t4`.
4833#[derive(Clone, Copy, Debug)]
4834pub struct TwoSeed<const K: usize> {
4835    /// The `ε⁰δ⁰` part: value / grad / Hessian of `ℓ`.
4836    pub base: Order2<K>,
4837    /// The `ε¹δ⁰` part.
4838    pub eps: Order2<K>,
4839    /// The `ε⁰δ¹` part.
4840    pub del: Order2<K>,
4841    /// The `ε¹δ¹` part. After a `seed(u, v)` evaluation,
4842    /// `eps_del.h[a][b] = Σ_{cd} ℓ_{abcd} u_c v_d`.
4843    pub eps_del: Order2<K>,
4844}
4845
4846impl<const K: usize> TwoSeed<K> {
4847    /// Seed primary `axis` at value `x` with ε-direction `u_axis` and
4848    /// δ-direction `v_axis`:
4849    /// `p_axis = p_axis⁰ + x-seed + ε·u_axis + δ·v_axis` (doc §A.3 "Seeding").
4850    pub fn seed(x: f64, axis: usize, u_axis: f64, v_axis: f64) -> Self {
4851        TwoSeed {
4852            base: Order2::variable(x, axis),
4853            eps: Order2::constant(u_axis),
4854            del: Order2::constant(v_axis),
4855            eps_del: Order2::constant(0.0),
4856        }
4857    }
4858
4859    /// The contracted-fourth channel after a `seed(u, v)` evaluation:
4860    /// `out[a][b] = Σ_{cd} ℓ_{abcd} u_c v_d`, i.e. the εδ-coefficient's Hessian.
4861    pub fn contracted_fourth(&self) -> [[f64; K]; K] {
4862        *self.eps_del.h()
4863    }
4864}
4865
4866impl<const K: usize> JetScalar<K> for TwoSeed<K> {
4867    fn constant(c: f64) -> Self {
4868        TwoSeed {
4869            base: Order2::constant(c),
4870            eps: Order2::constant(0.0),
4871            del: Order2::constant(0.0),
4872            eps_del: Order2::constant(0.0),
4873        }
4874    }
4875    fn variable(x: f64, axis: usize) -> Self {
4876        TwoSeed {
4877            base: Order2::variable(x, axis),
4878            eps: Order2::constant(0.0),
4879            del: Order2::constant(0.0),
4880            eps_del: Order2::constant(0.0),
4881        }
4882    }
4883}
4884
4885impl<const K: usize> crate::nested_dual::JetField for TwoSeed<K> {
4886    fn value(&self) -> f64 {
4887        self.base.value()
4888    }
4889    fn add(&self, o: &Self) -> Self {
4890        TwoSeed {
4891            base: self.base.add(&o.base),
4892            eps: self.eps.add(&o.eps),
4893            del: self.del.add(&o.del),
4894            eps_del: self.eps_del.add(&o.eps_del),
4895        }
4896    }
4897    fn sub(&self, o: &Self) -> Self {
4898        TwoSeed {
4899            base: self.base.sub(&o.base),
4900            eps: self.eps.sub(&o.eps),
4901            del: self.del.sub(&o.del),
4902            eps_del: self.eps_del.sub(&o.eps_del),
4903        }
4904    }
4905    fn mul(&self, o: &Self) -> Self {
4906        let a = self;
4907        let b = o;
4908        // Truncate ε² = δ² = 0 (doc §A.3 product table).
4909        let base = a.base.mul(&b.base);
4910        let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
4911        let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
4912        let eps_del = a
4913            .base
4914            .mul(&b.eps_del)
4915            .add(&a.eps.mul(&b.del))
4916            .add(&a.del.mul(&b.eps))
4917            .add(&a.eps_del.mul(&b.base));
4918        TwoSeed {
4919            base,
4920            eps,
4921            del,
4922            eps_del,
4923        }
4924    }
4925    fn neg(&self) -> Self {
4926        TwoSeed {
4927            base: self.base.neg(),
4928            eps: self.eps.neg(),
4929            del: self.del.neg(),
4930            eps_del: self.eps_del.neg(),
4931        }
4932    }
4933    fn scale(&self, s: f64) -> Self {
4934        TwoSeed {
4935            base: self.base.scale(s),
4936            eps: self.eps.scale(s),
4937            del: self.del.scale(s),
4938            eps_del: self.eps_del.scale(s),
4939        }
4940    }
4941    fn compose_unary(&self, d: [f64; 5]) -> Self {
4942        // f(s) with s = base + ε eps + δ del + εδ eps_del, ε²=δ²=0:
4943        //   f(s) = f(base)
4944        //        + ε · f'(base)·eps
4945        //        + δ · f'(base)·del
4946        //        + εδ · ( f''(base)·eps·del + f'(base)·eps_del ).
4947        // Each f^{(r)}(base) is the Order2 composition of base with the stack
4948        // shifted r entries (doc §A.3 composition). Order2 reads only the
4949        // leading three entries of whatever stack it is handed, so the trailing
4950        // padding slots are unused (the fixed-length array makes this total).
4951        let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
4952        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]); // f'(base) as Order2
4953        let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]); // f''(base) as Order2
4954        let eps = fprime.mul(&self.eps);
4955        let del = fprime.mul(&self.del);
4956        let eps_del = fsecond
4957            .mul(&self.eps)
4958            .mul(&self.del)
4959            .add(&fprime.mul(&self.eps_del));
4960        TwoSeed {
4961            base,
4962            eps,
4963            del,
4964            eps_del,
4965        }
4966    }
4967}
4968
4969// ── TwoSeedLane<L, K>: lane-batched two-seed, contracted fourth (doc §A.3) ─
4970
4971/// Lane-batched [`TwoSeed`]: the same two-seed scalar with its four [`Order2`]
4972/// parts re-typed to [`Order2Lane<L, K>`], so one `L = f64x4` instance carries
4973/// FOUR rows' contracted-fourth evaluations per vector pass.
4974///
4975/// Every operation is a term-for-term structural re-type of the scalar
4976/// [`TwoSeed`] ops onto the lane-implemented [`Order2Lane`] algebra. With
4977/// `L = f64`, `TwoSeedLane<f64, K>` is `to_bits`-identical to [`TwoSeed<K>`];
4978/// with `L = f64x4`, lane `i` is `to_bits`-identical to that (see `batch_tests`).
4979#[derive(Clone, Copy, Debug)]
4980pub struct TwoSeedLane<L: Lane, const K: usize> {
4981    /// The `ε⁰δ⁰` part.
4982    pub base: Order2Lane<L, K>,
4983    /// The `ε¹δ⁰` part.
4984    pub eps: Order2Lane<L, K>,
4985    /// The `ε⁰δ¹` part.
4986    pub del: Order2Lane<L, K>,
4987    /// The `ε¹δ¹` part. After a `seed(u, v)` evaluation, `eps_del.h[a][b]`
4988    /// lane `i` is row `i`'s `Σ_{cd} ℓ_{abcd} u_c v_d`.
4989    pub eps_del: Order2Lane<L, K>,
4990}
4991
4992/// The 4-rows-per-pass batched two-seed scalar (`wide::f64x4` lanes).
4993pub type TwoSeedBatch<const K: usize> = TwoSeedLane<wide::f64x4, K>;
4994
4995impl<L: Lane, const K: usize> TwoSeedLane<L, K> {
4996    /// A constant: base = `constant(c)`, all seed parts zero (mirrors
4997    /// [`TwoSeed::constant`]).
4998    #[inline]
4999    pub fn constant(c: L) -> Self {
5000        let z = Order2Lane::constant(L::splat(0.0));
5001        TwoSeedLane {
5002            base: Order2Lane::constant(c),
5003            eps: z,
5004            del: z,
5005            eps_del: z,
5006        }
5007    }
5008
5009    /// The seeded variable `p_axis` at (per-lane) value `value`, no ε/δ direction
5010    /// (mirrors [`TwoSeed::variable`]).
5011    #[inline]
5012    pub fn variable(value: L, axis: usize) -> Self {
5013        let z = Order2Lane::constant(L::splat(0.0));
5014        TwoSeedLane {
5015            base: Order2Lane::variable(value, axis),
5016            eps: z,
5017            del: z,
5018            eps_del: z,
5019        }
5020    }
5021
5022    /// Seed primary `axis` at (per-lane) value `value` with ε-direction `u_axis`
5023    /// and δ-direction `v_axis` (mirrors [`TwoSeed::seed`]). With `L = f64x4`,
5024    /// each argument packs the four rows' values for primary `axis`.
5025    #[inline]
5026    pub fn seed(value: L, axis: usize, u_axis: L, v_axis: L) -> Self {
5027        TwoSeedLane {
5028            base: Order2Lane::variable(value, axis),
5029            eps: Order2Lane::constant(u_axis),
5030            del: Order2Lane::constant(v_axis),
5031            eps_del: Order2Lane::constant(L::splat(0.0)),
5032        }
5033    }
5034
5035    /// The contracted-fourth channel after a `seed(u, v)` evaluation:
5036    /// `out[a][b]` lane `i` is row `i`'s `Σ_{cd} ℓ_{abcd} u_c v_d`
5037    /// (the εδ-part Hessian).
5038    #[inline]
5039    #[must_use]
5040    pub fn contracted_fourth(&self) -> [[L; K]; K] {
5041        self.eps_del.h
5042    }
5043
5044    /// Lane-wise `self + o` (mirrors [`TwoSeed::add`]).
5045    #[inline]
5046    pub fn add(&self, o: &Self) -> Self {
5047        TwoSeedLane {
5048            base: self.base.add(&o.base),
5049            eps: self.eps.add(&o.eps),
5050            del: self.del.add(&o.del),
5051            eps_del: self.eps_del.add(&o.eps_del),
5052        }
5053    }
5054
5055    /// Lane-wise `self - o` (mirrors [`TwoSeed::sub`]).
5056    #[inline]
5057    pub fn sub(&self, o: &Self) -> Self {
5058        TwoSeedLane {
5059            base: self.base.sub(&o.base),
5060            eps: self.eps.sub(&o.eps),
5061            del: self.del.sub(&o.del),
5062            eps_del: self.eps_del.sub(&o.eps_del),
5063        }
5064    }
5065
5066    /// Lane-wise `self · o`, ε² = δ² = 0 truncation (mirrors [`TwoSeed::mul`]).
5067    #[inline]
5068    pub fn mul(&self, o: &Self) -> Self {
5069        let a = self;
5070        let b = o;
5071        let base = a.base.mul(&b.base);
5072        let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
5073        let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
5074        let eps_del = a
5075            .base
5076            .mul(&b.eps_del)
5077            .add(&a.eps.mul(&b.del))
5078            .add(&a.del.mul(&b.eps))
5079            .add(&a.eps_del.mul(&b.base));
5080        TwoSeedLane {
5081            base,
5082            eps,
5083            del,
5084            eps_del,
5085        }
5086    }
5087
5088    /// Negate every part (mirrors [`TwoSeed::neg`]).
5089    #[inline]
5090    pub fn neg(&self) -> Self {
5091        TwoSeedLane {
5092            base: self.base.neg(),
5093            eps: self.eps.neg(),
5094            del: self.del.neg(),
5095            eps_del: self.eps_del.neg(),
5096        }
5097    }
5098
5099    /// Multiply every part by the plain scalar `s` (mirrors [`TwoSeed::scale`]).
5100    #[inline]
5101    pub fn scale(&self, s: f64) -> Self {
5102        TwoSeedLane {
5103            base: self.base.scale(s),
5104            eps: self.eps.scale(s),
5105            del: self.del.scale(s),
5106            eps_del: self.eps_del.scale(s),
5107        }
5108    }
5109
5110    /// Exact composition `f ∘ self`, given the per-lane outer-derivative stack
5111    /// `d = [f, f′, f″, f‴, f⁗]`. Term-for-term identical to
5112    /// [`TwoSeed::compose_unary`]: base reads `d[0..=2]`, `f′(base)` reads
5113    /// `d[1..=3]`, `f″(base)` reads `d[2..=4]`, and the cross part carries
5114    /// `f″·eps·del + f′·eps_del`.
5115    #[inline]
5116    pub fn compose_unary(&self, d: [L; 5]) -> Self {
5117        let base = self.base.compose_unary([d[0], d[1], d[2]]);
5118        let fprime = self.base.compose_unary([d[1], d[2], d[3]]);
5119        let fsecond = self.base.compose_unary([d[2], d[3], d[4]]);
5120        let eps = fprime.mul(&self.eps);
5121        let del = fprime.mul(&self.del);
5122        let eps_del = fsecond
5123            .mul(&self.eps)
5124            .mul(&self.del)
5125            .add(&fprime.mul(&self.eps_del));
5126        TwoSeedLane {
5127            base,
5128            eps,
5129            del,
5130            eps_del,
5131        }
5132    }
5133
5134    /// `e^self`, per-lane stack `[e; 5]` (matches [`JetScalar::exp`]).
5135    #[inline]
5136    pub fn exp(&self) -> Self {
5137        let d = self.base.v.unary5(|u| {
5138            let e = u.exp();
5139            [e, e, e, e, e]
5140        });
5141        self.compose_unary(d)
5142    }
5143
5144    /// `ln(self)`; caller guarantees positivity (matches [`JetScalar::ln`]).
5145    #[inline]
5146    pub fn ln(&self) -> Self {
5147        let d = self.base.v.unary5(|u| {
5148            let r = 1.0 / u;
5149            [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
5150        });
5151        self.compose_unary(d)
5152    }
5153
5154    /// `√self`; caller guarantees positivity (matches [`JetScalar::sqrt`]).
5155    #[inline]
5156    pub fn sqrt(&self) -> Self {
5157        let d = self.base.v.unary5(|u| {
5158            let s = u.sqrt();
5159            [
5160                s,
5161                0.5 / s,
5162                -0.25 / (u * s),
5163                0.375 / (u * u * s),
5164                -0.9375 / (u * u * u * s),
5165            ]
5166        });
5167        self.compose_unary(d)
5168    }
5169
5170    /// `1/self` (matches [`JetScalar::recip`]).
5171    #[inline]
5172    pub fn recip(&self) -> Self {
5173        let d = self.base.v.unary5(|u| {
5174            let r = 1.0 / u;
5175            let r2 = r * r;
5176            [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
5177        });
5178        self.compose_unary(d)
5179    }
5180
5181    /// `self^a` for real `a`; caller guarantees a positive base (matches
5182    /// [`JetScalar::powf`]).
5183    #[inline]
5184    pub fn powf(&self, a: f64) -> Self {
5185        let d = self.base.v.unary5(|u| {
5186            [
5187                u.powf(a),
5188                a * u.powf(a - 1.0),
5189                a * (a - 1.0) * u.powf(a - 2.0),
5190                a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
5191                a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
5192            ]
5193        });
5194        self.compose_unary(d)
5195    }
5196
5197    /// `ln Γ(self)`; caller guarantees positivity (matches [`JetScalar::ln_gamma`]).
5198    #[inline]
5199    pub fn ln_gamma(&self) -> Self {
5200        let d = self
5201            .base
5202            .v
5203            .unary5(crate::jet_tower::ln_gamma_derivative_stack);
5204        self.compose_unary(d)
5205    }
5206
5207    /// `ψ(self)` digamma; caller guarantees positivity (matches
5208    /// [`JetScalar::digamma`]).
5209    #[inline]
5210    pub fn digamma(&self) -> Self {
5211        let d = self
5212            .base
5213            .v
5214            .unary5(crate::jet_tower::digamma_derivative_stack);
5215        self.compose_unary(d)
5216    }
5217}
5218
5219impl<const K: usize> TwoSeedBatch<K> {
5220    /// Extract lane `i`'s parts as a production [`TwoSeed<K>`]. Lane `i` is
5221    /// `to_bits`-identical to evaluating the same program at [`TwoSeed<K>`] on
5222    /// row `i` (see `batch_tests`).
5223    #[inline]
5224    #[must_use]
5225    pub fn lane(&self, i: usize) -> TwoSeed<K> {
5226        TwoSeed {
5227            base: self.base.lane(i),
5228            eps: self.eps.lane(i),
5229            del: self.del.lane(i),
5230            eps_del: self.eps_del.lane(i),
5231        }
5232    }
5233}
5234
5235// ── Tower3<K>: value / gradient / Hessian / third tensor ────────────────
5236
5237/// The order-≤3 [`crate::jet_tower::Tower3`] is also a [`JetScalar`]. It serves
5238/// consumers that read `.t3` but never `.t4`, avoiding the fourth-tensor
5239/// product/composition work while preserving the lower channels
5240/// bit-for-bit against [`crate::jet_tower::Tower4`].
5241impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower3<K> {
5242    fn constant(c: f64) -> Self {
5243        crate::jet_tower::Tower3::constant(c)
5244    }
5245    fn variable(x: f64, axis: usize) -> Self {
5246        crate::jet_tower::Tower3::variable(x, axis)
5247    }
5248}
5249
5250impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower3<K> {
5251    fn value(&self) -> f64 {
5252        self.v
5253    }
5254    fn add(&self, o: &Self) -> Self {
5255        *self + *o
5256    }
5257    fn sub(&self, o: &Self) -> Self {
5258        *self + o.scale(-1.0)
5259    }
5260    fn mul(&self, o: &Self) -> Self {
5261        crate::jet_tower::Tower3::mul(self, o)
5262    }
5263    fn neg(&self) -> Self {
5264        self.scale(-1.0)
5265    }
5266    fn scale(&self, s: f64) -> Self {
5267        crate::jet_tower::Tower3::scale(self, s)
5268    }
5269    fn compose_unary(&self, d: [f64; 5]) -> Self {
5270        crate::jet_tower::Tower3::compose_unary(self, [d[0], d[1], d[2], d[3]])
5271    }
5272}
5273
5274// ── Tower4<K>: full dense tower as a JetScalar (the all-channels scalar) ─
5275
5276/// The full dense [`crate::jet_tower::Tower4`] is itself a [`JetScalar`]: it
5277/// carries EVERY channel, so a row expression written ONCE against [`JetScalar`]
5278/// can be evaluated at `Tower4` to obtain the full `(v, g, H, t3, t4)` in one
5279/// pass. This is BOTH the #932 oracle ground truth the packed [`Order2`] /
5280/// [`OneSeed`] / [`TwoSeed`] scalars are pinned against, AND a production scalar:
5281/// a family whose uncontracted third / fourth derivative tensors are needed
5282/// (the BMS rigid `third_full` / `fourth_full` caches) evaluates the SAME
5283/// generic row-NLL expression at `Tower4` and reads `.t3` / `.t4` off the
5284/// result — so the dense tensors come from the single source of truth, not a
5285/// separately hand-written jet. The packed scalars serve the consumers that
5286/// need only `(v, g, H)` (`Order2`) or one / two contractions
5287/// (`OneSeed` / `TwoSeed`) without paying for the dense tensors.
5288impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower4<K> {
5289    fn constant(c: f64) -> Self {
5290        crate::jet_tower::Tower4::constant(c)
5291    }
5292    fn variable(x: f64, axis: usize) -> Self {
5293        crate::jet_tower::Tower4::variable(x, axis)
5294    }
5295}
5296
5297impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower4<K> {
5298    fn value(&self) -> f64 {
5299        self.v
5300    }
5301    fn add(&self, o: &Self) -> Self {
5302        *self + *o
5303    }
5304    fn sub(&self, o: &Self) -> Self {
5305        *self - *o
5306    }
5307    fn mul(&self, o: &Self) -> Self {
5308        crate::jet_tower::Tower4::mul(self, o)
5309    }
5310    fn neg(&self) -> Self {
5311        self.scale(-1.0)
5312    }
5313    fn scale(&self, s: f64) -> Self {
5314        crate::jet_tower::Tower4::scale(self, s)
5315    }
5316    fn compose_unary(&self, d: [f64; 5]) -> Self {
5317        crate::jet_tower::Tower4::compose_unary(self, d)
5318    }
5319}
5320
5321#[cfg(test)]
5322mod tests {
5323    use super::*;
5324    use crate::jet_tower::{RowProgram, Tower4, program_full_tower};
5325    use crate::nested_dual::JetField;
5326
5327    struct DenseSymmetric3([[f64; 3]; 3]);
5328
5329    impl SymmetricQuadraticCoefficients for DenseSymmetric3 {
5330        fn dimension(&self) -> usize {
5331            3
5332        }
5333
5334        fn multiply(&self, input: &[f64], output: &mut [f64]) {
5335            assert_eq!(input.len(), 3);
5336            assert_eq!(output.len(), 3);
5337            for (row, output) in output.iter_mut().enumerate() {
5338                *output = (0..3)
5339                    .map(|column| self.0[row][column] * input[column])
5340                    .sum();
5341            }
5342        }
5343
5344        fn coefficient(&self, row: usize, column: usize) -> f64 {
5345            self.0[row][column]
5346        }
5347    }
5348
5349    #[test]
5350    fn symmetric_quadratic_order2_lowerings_match_scalar_program() {
5351        const K: usize = 4;
5352        let coefficients = DenseSymmetric3([[1.2, 0.3, -0.2], [0.3, 0.8, 0.15], [-0.2, 0.15, 1.5]]);
5353        let values = [0.4, -0.7, 1.1, 0.25];
5354        let fixed_vars: [Order2<K>; K] =
5355            std::array::from_fn(|axis| Order2::variable(values[axis], axis));
5356        let fixed_inputs = [
5357            fixed_vars[0].mul(&fixed_vars[1]).add(&fixed_vars[3]),
5358            fixed_vars[1].exp().add(&fixed_vars[2].scale(0.4)),
5359            fixed_vars[2].mul(&fixed_vars[2]).sub(&fixed_vars[0]),
5360        ];
5361        let fixed_direct = Order2::symmetric_quadratic_form(&fixed_inputs, &coefficients);
5362        let fixed_scalar = symmetric_quadratic_form_default(
5363            &fixed_inputs,
5364            &coefficients,
5365            Order2::constant,
5366            JetField::add,
5367            JetField::mul,
5368            JetField::scale,
5369        );
5370        let weights = [0.7, -1.1, 0.35];
5371        let fixed_linear_direct = Order2::linear_combination(&fixed_inputs, &weights);
5372        let fixed_linear_scalar = linear_combination_default(
5373            &fixed_inputs,
5374            &weights,
5375            Order2::constant,
5376            JetField::add,
5377            JetField::scale,
5378        );
5379        let derivative_stacks = [
5380            [0.8, -0.3, 0.7, 0.0, 0.0],
5381            [-0.2, 1.1, -0.4, 0.0, 0.0],
5382            [1.4, 0.25, 0.6, 0.0, 0.0],
5383        ];
5384        let fixed_add_direct = fixed_inputs[0].add_constant(0.65);
5385        let fixed_add_scalar = fixed_inputs[0].add(&Order2::constant(0.65));
5386        let fixed_multiply_add_direct =
5387            fixed_inputs[0].multiply_add(&fixed_inputs[1], &fixed_inputs[2]);
5388        let fixed_multiply_add_scalar = multiply_add_default(
5389            &fixed_inputs[0],
5390            &fixed_inputs[1],
5391            &fixed_inputs[2],
5392            JetField::mul,
5393            JetField::add,
5394        );
5395        let fixed_composed_direct = Order2::composed_sum(&fixed_inputs, &derivative_stacks);
5396        let fixed_composed_scalar = composed_sum_default(
5397            &fixed_inputs,
5398            &derivative_stacks,
5399            Order2::constant,
5400            JetField::add,
5401            JetField::compose_unary,
5402        );
5403
5404        let value_vars: [RuntimeValue; K] =
5405            std::array::from_fn(|axis| RuntimeValue::variable(values[axis], axis, K, &()));
5406        let value_inputs = [
5407            value_vars[0].mul(&value_vars[1]).add(&value_vars[3]),
5408            value_vars[1].exp().add(&value_vars[2].scale(0.4)),
5409            value_vars[2].mul(&value_vars[2]).sub(&value_vars[0]),
5410        ];
5411        let value_quadratic =
5412            RuntimeValue::symmetric_quadratic_form(&value_inputs, &coefficients, K, &());
5413        let value_linear = RuntimeValue::linear_combination(&value_inputs, &weights, K, &());
5414        let value_composed = RuntimeValue::composed_sum(&value_inputs, &derivative_stacks, K, &());
5415
5416        let arena = DynamicJetArena::new();
5417        let dynamic_vars: [DynamicOrder2<'_>; K] =
5418            std::array::from_fn(|axis| DynamicOrder2::variable(values[axis], axis, K, &arena));
5419        let dynamic_inputs = [
5420            dynamic_vars[0].mul(&dynamic_vars[1]).add(&dynamic_vars[3]),
5421            dynamic_vars[1].exp().add(&dynamic_vars[2].scale(0.4)),
5422            dynamic_vars[2].mul(&dynamic_vars[2]).sub(&dynamic_vars[0]),
5423        ];
5424        let dynamic_direct =
5425            DynamicOrder2::symmetric_quadratic_form(&dynamic_inputs, &coefficients, K, &arena);
5426        let dynamic_scalar = symmetric_quadratic_form_default(
5427            &dynamic_inputs,
5428            &coefficients,
5429            |value| DynamicOrder2::constant(value, K, &arena),
5430            RuntimeJetScalar::add,
5431            RuntimeJetScalar::mul,
5432            RuntimeJetScalar::scale,
5433        );
5434        let dynamic_linear_direct =
5435            DynamicOrder2::linear_combination(&dynamic_inputs, &weights, K, &arena);
5436        let dynamic_linear_scalar = linear_combination_default(
5437            &dynamic_inputs,
5438            &weights,
5439            |value| DynamicOrder2::constant(value, K, &arena),
5440            RuntimeJetScalar::add,
5441            RuntimeJetScalar::scale,
5442        );
5443        let dynamic_add_direct = dynamic_inputs[0].add_constant(0.65, &arena);
5444        let dynamic_add_scalar = dynamic_inputs[0].add(&DynamicOrder2::constant(0.65, K, &arena));
5445        let dynamic_multiply_add_direct =
5446            dynamic_inputs[0].multiply_add(&dynamic_inputs[1], &dynamic_inputs[2]);
5447        let dynamic_multiply_add_scalar = multiply_add_default(
5448            &dynamic_inputs[0],
5449            &dynamic_inputs[1],
5450            &dynamic_inputs[2],
5451            RuntimeJetScalar::mul,
5452            RuntimeJetScalar::add,
5453        );
5454        let dynamic_composed_direct =
5455            DynamicOrder2::composed_sum(&dynamic_inputs, &derivative_stacks, K, &arena);
5456        let dynamic_composed_scalar = composed_sum_default(
5457            &dynamic_inputs,
5458            &derivative_stacks,
5459            |value| DynamicOrder2::constant(value, K, &arena),
5460            RuntimeJetScalar::add,
5461            RuntimeJetScalar::compose_unary,
5462        );
5463
5464        let tolerance = 2.0e-13;
5465        for (label, actual, expected) in [
5466            ("fixed value", fixed_direct.value(), fixed_scalar.value()),
5467            (
5468                "zero-order quadratic value",
5469                value_quadratic.value(),
5470                fixed_direct.value(),
5471            ),
5472            (
5473                "zero-order linear value",
5474                value_linear.value(),
5475                fixed_linear_direct.value(),
5476            ),
5477            (
5478                "zero-order composed value",
5479                value_composed.value(),
5480                fixed_composed_direct.value(),
5481            ),
5482            (
5483                "dynamic value",
5484                dynamic_direct.value(),
5485                dynamic_scalar.value(),
5486            ),
5487            (
5488                "fixed linear value",
5489                fixed_linear_direct.value(),
5490                fixed_linear_scalar.value(),
5491            ),
5492            (
5493                "dynamic linear value",
5494                dynamic_linear_direct.value(),
5495                dynamic_linear_scalar.value(),
5496            ),
5497            (
5498                "fixed add-constant value",
5499                fixed_add_direct.value(),
5500                fixed_add_scalar.value(),
5501            ),
5502            (
5503                "dynamic add-constant value",
5504                dynamic_add_direct.value(),
5505                dynamic_add_scalar.value(),
5506            ),
5507            (
5508                "fixed multiply-add value",
5509                fixed_multiply_add_direct.value(),
5510                fixed_multiply_add_scalar.value(),
5511            ),
5512            (
5513                "dynamic multiply-add value",
5514                dynamic_multiply_add_direct.value(),
5515                dynamic_multiply_add_scalar.value(),
5516            ),
5517            (
5518                "fixed composed-sum value",
5519                fixed_composed_direct.value(),
5520                fixed_composed_scalar.value(),
5521            ),
5522            (
5523                "dynamic composed-sum value",
5524                dynamic_composed_direct.value(),
5525                dynamic_composed_scalar.value(),
5526            ),
5527        ] {
5528            assert!(
5529                (actual - expected).abs() <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5530                "{label}: direct={actual:+.16e}, scalar={expected:+.16e}"
5531            );
5532        }
5533        for primary_a in 0..K {
5534            for (label, actual, expected) in [
5535                (
5536                    "fixed gradient",
5537                    fixed_direct.g()[primary_a],
5538                    fixed_scalar.g()[primary_a],
5539                ),
5540                (
5541                    "dynamic gradient",
5542                    dynamic_direct.g()[primary_a],
5543                    dynamic_scalar.g()[primary_a],
5544                ),
5545                (
5546                    "fixed linear gradient",
5547                    fixed_linear_direct.g()[primary_a],
5548                    fixed_linear_scalar.g()[primary_a],
5549                ),
5550                (
5551                    "dynamic linear gradient",
5552                    dynamic_linear_direct.g()[primary_a],
5553                    dynamic_linear_scalar.g()[primary_a],
5554                ),
5555                (
5556                    "fixed add-constant gradient",
5557                    fixed_add_direct.g()[primary_a],
5558                    fixed_add_scalar.g()[primary_a],
5559                ),
5560                (
5561                    "dynamic add-constant gradient",
5562                    dynamic_add_direct.g()[primary_a],
5563                    dynamic_add_scalar.g()[primary_a],
5564                ),
5565                (
5566                    "fixed multiply-add gradient",
5567                    fixed_multiply_add_direct.g()[primary_a],
5568                    fixed_multiply_add_scalar.g()[primary_a],
5569                ),
5570                (
5571                    "dynamic multiply-add gradient",
5572                    dynamic_multiply_add_direct.g()[primary_a],
5573                    dynamic_multiply_add_scalar.g()[primary_a],
5574                ),
5575                (
5576                    "fixed composed-sum gradient",
5577                    fixed_composed_direct.g()[primary_a],
5578                    fixed_composed_scalar.g()[primary_a],
5579                ),
5580                (
5581                    "dynamic composed-sum gradient",
5582                    dynamic_composed_direct.g()[primary_a],
5583                    dynamic_composed_scalar.g()[primary_a],
5584                ),
5585            ] {
5586                assert!(
5587                    (actual - expected).abs()
5588                        <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5589                    "{label}[{primary_a}]: direct={actual:+.16e}, scalar={expected:+.16e}"
5590                );
5591            }
5592            for primary_b in 0..K {
5593                for (label, actual, expected) in [
5594                    (
5595                        "fixed Hessian",
5596                        fixed_direct.h()[primary_a][primary_b],
5597                        fixed_scalar.h()[primary_a][primary_b],
5598                    ),
5599                    (
5600                        "dynamic Hessian",
5601                        dynamic_direct.h_at(primary_a, primary_b),
5602                        dynamic_scalar.h_at(primary_a, primary_b),
5603                    ),
5604                    (
5605                        "fixed linear Hessian",
5606                        fixed_linear_direct.h()[primary_a][primary_b],
5607                        fixed_linear_scalar.h()[primary_a][primary_b],
5608                    ),
5609                    (
5610                        "dynamic linear Hessian",
5611                        dynamic_linear_direct.h_at(primary_a, primary_b),
5612                        dynamic_linear_scalar.h_at(primary_a, primary_b),
5613                    ),
5614                    (
5615                        "fixed add-constant Hessian",
5616                        fixed_add_direct.h()[primary_a][primary_b],
5617                        fixed_add_scalar.h()[primary_a][primary_b],
5618                    ),
5619                    (
5620                        "dynamic add-constant Hessian",
5621                        dynamic_add_direct.h_at(primary_a, primary_b),
5622                        dynamic_add_scalar.h_at(primary_a, primary_b),
5623                    ),
5624                    (
5625                        "fixed multiply-add Hessian",
5626                        fixed_multiply_add_direct.h()[primary_a][primary_b],
5627                        fixed_multiply_add_scalar.h()[primary_a][primary_b],
5628                    ),
5629                    (
5630                        "dynamic multiply-add Hessian",
5631                        dynamic_multiply_add_direct.h_at(primary_a, primary_b),
5632                        dynamic_multiply_add_scalar.h_at(primary_a, primary_b),
5633                    ),
5634                    (
5635                        "fixed composed-sum Hessian",
5636                        fixed_composed_direct.h()[primary_a][primary_b],
5637                        fixed_composed_scalar.h()[primary_a][primary_b],
5638                    ),
5639                    (
5640                        "dynamic composed-sum Hessian",
5641                        dynamic_composed_direct.h_at(primary_a, primary_b),
5642                        dynamic_composed_scalar.h_at(primary_a, primary_b),
5643                    ),
5644                ] {
5645                    assert!(
5646                        (actual - expected).abs()
5647                            <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5648                        "{label}[{primary_a},{primary_b}]: direct={actual:+.16e}, scalar={expected:+.16e}"
5649                    );
5650                }
5651            }
5652        }
5653    }
5654
5655    #[test]
5656    fn compiled_product_affine_and_fused_nodes_match_scalar_program_randomized() {
5657        const K: usize = 4;
5658        const TERMS: usize = 10;
5659
5660        fn sample(state: &mut u64) -> f64 {
5661            *state ^= *state << 13;
5662            *state ^= *state >> 7;
5663            *state ^= *state << 17;
5664            let unit = (*state >> 11) as f64 * (1.0 / ((1_u64 << 53) as f64));
5665            2.0 * unit - 1.0
5666        }
5667
5668        fn arbitrary_order2<const K: usize>(state: &mut u64) -> Order2<K> {
5669            let mut tower = crate::jet_tower::Tower2::zero();
5670            tower.v = sample(state);
5671            for primary in 0..K {
5672                tower.g[primary] = sample(state);
5673                for other in primary..K {
5674                    let channel = sample(state);
5675                    tower.h[primary][other] = channel;
5676                    tower.h[other][primary] = channel;
5677                }
5678            }
5679            Order2(tower)
5680        }
5681
5682        fn close(actual: f64, expected: f64, case: usize, label: &str) {
5683            let tolerance = 2.0e-12 * actual.abs().max(expected.abs()).max(1.0);
5684            assert!(
5685                (actual - expected).abs() <= tolerance,
5686                "case {case} {label}: direct={actual:+.16e}, scalar={expected:+.16e}, tolerance={tolerance:.3e}"
5687            );
5688        }
5689
5690        let mut state = 0x932a_ff1e_c0de_5eed_u64;
5691        for case in 0..256 {
5692            // Arbitrary value, gradient, and symmetric-Hessian inputs model the
5693            // complete local state of nonlinear upstream expressions. This is
5694            // stronger than testing only independently seeded variables.
5695            let fixed_inputs: [Order2<K>; TERMS] =
5696                std::array::from_fn(|_| arbitrary_order2(&mut state));
5697            let mut input_scales: [f64; TERMS] = std::array::from_fn(|_| sample(&mut state));
5698            input_scales[0] = -1.25;
5699            input_scales[1] = 0.0;
5700            input_scales[4] = 1.25;
5701            let addend_scales: [f64; TERMS] = std::array::from_fn(|term| match term % 4 {
5702                0 => 0.0,
5703                1 => 1.0,
5704                2 => -0.75,
5705                _ => 0.35,
5706            });
5707            let derivative_stacks: [[f64; 5]; TERMS] =
5708                std::array::from_fn(|_| std::array::from_fn(|_| sample(&mut state)));
5709            let input_shift = sample(&mut state);
5710            let mut fixed_lefts = std::array::from_fn(|term| &fixed_inputs[term]);
5711            fixed_lefts[4] = &fixed_inputs[0];
5712            fixed_lefts[5] = &fixed_inputs[0];
5713            let fixed_right = &fixed_inputs[1];
5714            let fixed_addend = &fixed_inputs[2];
5715
5716            let fixed_product_direct = fixed_inputs[0].product(&fixed_inputs[1]);
5717            let fixed_product_scalar = fixed_inputs[0].mul(&fixed_inputs[1]);
5718            let fixed_affine_direct =
5719                fixed_inputs[2].affine_compose(input_scales[2], input_shift, derivative_stacks[2]);
5720            let fixed_affine_scalar = affine_compose_default(
5721                &fixed_inputs[2],
5722                input_scales[2],
5723                input_shift,
5724                derivative_stacks[2],
5725                JetField::scale,
5726                Order2::add_constant,
5727                JetField::compose_unary,
5728            );
5729            let fixed_sum_direct =
5730                Order2::affine_composed_sum(&fixed_inputs, &input_scales, &derivative_stacks);
5731            let fixed_sum_scalar = affine_composed_sum_default(
5732                &fixed_inputs,
5733                &input_scales,
5734                &derivative_stacks,
5735                Order2::constant,
5736                JetField::add,
5737                JetField::scale,
5738                Order2::add_constant,
5739                JetField::compose_unary,
5740            );
5741            let fixed_fused_direct = Order2::shared_multiply_add_affine_composed_sum(
5742                &fixed_lefts,
5743                fixed_right,
5744                fixed_addend,
5745                &addend_scales,
5746                &input_scales,
5747                &derivative_stacks,
5748            );
5749            let fixed_fused_scalar = shared_multiply_add_affine_composed_sum_default(
5750                &fixed_lefts,
5751                fixed_right,
5752                fixed_addend,
5753                &addend_scales,
5754                &input_scales,
5755                &derivative_stacks,
5756                Order2::constant,
5757                JetField::add,
5758                JetField::mul,
5759                JetField::scale,
5760                Order2::multiply_add,
5761                Order2::affine_compose,
5762            );
5763
5764            let arena = DynamicJetArena::new();
5765            let dynamic_inputs: [DynamicOrder2<'_>; TERMS] = std::array::from_fn(|term| {
5766                DynamicOrder2::from_channel_functions(
5767                    fixed_inputs[term].value(),
5768                    K,
5769                    &arena,
5770                    |primary| fixed_inputs[term].g()[primary],
5771                    |primary, other| fixed_inputs[term].h()[primary][other],
5772                )
5773            });
5774            let dynamic_product_direct = dynamic_inputs[0].product(&dynamic_inputs[1]);
5775            let dynamic_product_scalar = dynamic_inputs[0].mul(&dynamic_inputs[1]);
5776            let dynamic_affine_direct = dynamic_inputs[2].affine_compose(
5777                input_scales[2],
5778                input_shift,
5779                derivative_stacks[2],
5780                &arena,
5781            );
5782            let dynamic_affine_scalar = affine_compose_default(
5783                &dynamic_inputs[2],
5784                input_scales[2],
5785                input_shift,
5786                derivative_stacks[2],
5787                RuntimeJetScalar::scale,
5788                |input, constant| input.add_constant(constant, &arena),
5789                RuntimeJetScalar::compose_unary,
5790            );
5791            let dynamic_sum_direct = DynamicOrder2::affine_composed_sum(
5792                &dynamic_inputs,
5793                &input_scales,
5794                &derivative_stacks,
5795                K,
5796                &arena,
5797            );
5798            let dynamic_sum_scalar = affine_composed_sum_default(
5799                &dynamic_inputs,
5800                &input_scales,
5801                &derivative_stacks,
5802                |value| DynamicOrder2::constant(value, K, &arena),
5803                RuntimeJetScalar::add,
5804                RuntimeJetScalar::scale,
5805                |input, constant| input.add_constant(constant, &arena),
5806                RuntimeJetScalar::compose_unary,
5807            );
5808            let mut dynamic_lefts = std::array::from_fn(|term| &dynamic_inputs[term]);
5809            dynamic_lefts[4] = &dynamic_inputs[0];
5810            dynamic_lefts[5] = &dynamic_inputs[0];
5811            let dynamic_right = &dynamic_inputs[1];
5812            let dynamic_addend = &dynamic_inputs[2];
5813            let dynamic_fused_direct = DynamicOrder2::shared_multiply_add_affine_composed_sum(
5814                &dynamic_lefts,
5815                dynamic_right,
5816                dynamic_addend,
5817                &addend_scales,
5818                &input_scales,
5819                &derivative_stacks,
5820                K,
5821                &arena,
5822            );
5823            let dynamic_fused_scalar = shared_multiply_add_affine_composed_sum_default(
5824                &dynamic_lefts,
5825                dynamic_right,
5826                dynamic_addend,
5827                &addend_scales,
5828                &input_scales,
5829                &derivative_stacks,
5830                |value| DynamicOrder2::constant(value, K, &arena),
5831                RuntimeJetScalar::add,
5832                RuntimeJetScalar::mul,
5833                RuntimeJetScalar::scale,
5834                RuntimeJetScalar::multiply_add,
5835                |input, scale, shift, stack| input.affine_compose(scale, shift, stack, &arena),
5836            );
5837
5838            for (label, actual, expected) in [
5839                (
5840                    "fixed product value",
5841                    fixed_product_direct.value(),
5842                    fixed_product_scalar.value(),
5843                ),
5844                (
5845                    "fixed affine value",
5846                    fixed_affine_direct.value(),
5847                    fixed_affine_scalar.value(),
5848                ),
5849                (
5850                    "fixed affine sum value",
5851                    fixed_sum_direct.value(),
5852                    fixed_sum_scalar.value(),
5853                ),
5854                (
5855                    "fixed fused value",
5856                    fixed_fused_direct.value(),
5857                    fixed_fused_scalar.value(),
5858                ),
5859                (
5860                    "dynamic product value",
5861                    dynamic_product_direct.value(),
5862                    dynamic_product_scalar.value(),
5863                ),
5864                (
5865                    "dynamic affine value",
5866                    dynamic_affine_direct.value(),
5867                    dynamic_affine_scalar.value(),
5868                ),
5869                (
5870                    "dynamic affine sum value",
5871                    dynamic_sum_direct.value(),
5872                    dynamic_sum_scalar.value(),
5873                ),
5874                (
5875                    "dynamic fused value",
5876                    dynamic_fused_direct.value(),
5877                    dynamic_fused_scalar.value(),
5878                ),
5879            ] {
5880                close(actual, expected, case, label);
5881            }
5882            for primary in 0..K {
5883                for (label, actual, expected) in [
5884                    (
5885                        "fixed product gradient",
5886                        fixed_product_direct.g()[primary],
5887                        fixed_product_scalar.g()[primary],
5888                    ),
5889                    (
5890                        "fixed affine gradient",
5891                        fixed_affine_direct.g()[primary],
5892                        fixed_affine_scalar.g()[primary],
5893                    ),
5894                    (
5895                        "fixed affine sum gradient",
5896                        fixed_sum_direct.g()[primary],
5897                        fixed_sum_scalar.g()[primary],
5898                    ),
5899                    (
5900                        "fixed fused gradient",
5901                        fixed_fused_direct.g()[primary],
5902                        fixed_fused_scalar.g()[primary],
5903                    ),
5904                    (
5905                        "dynamic product gradient",
5906                        dynamic_product_direct.g()[primary],
5907                        dynamic_product_scalar.g()[primary],
5908                    ),
5909                    (
5910                        "dynamic affine gradient",
5911                        dynamic_affine_direct.g()[primary],
5912                        dynamic_affine_scalar.g()[primary],
5913                    ),
5914                    (
5915                        "dynamic affine sum gradient",
5916                        dynamic_sum_direct.g()[primary],
5917                        dynamic_sum_scalar.g()[primary],
5918                    ),
5919                    (
5920                        "dynamic fused gradient",
5921                        dynamic_fused_direct.g()[primary],
5922                        dynamic_fused_scalar.g()[primary],
5923                    ),
5924                ] {
5925                    close(actual, expected, case, label);
5926                }
5927                for other in 0..K {
5928                    for (label, actual, expected) in [
5929                        (
5930                            "fixed product Hessian",
5931                            fixed_product_direct.h()[primary][other],
5932                            fixed_product_scalar.h()[primary][other],
5933                        ),
5934                        (
5935                            "fixed affine Hessian",
5936                            fixed_affine_direct.h()[primary][other],
5937                            fixed_affine_scalar.h()[primary][other],
5938                        ),
5939                        (
5940                            "fixed affine sum Hessian",
5941                            fixed_sum_direct.h()[primary][other],
5942                            fixed_sum_scalar.h()[primary][other],
5943                        ),
5944                        (
5945                            "fixed fused Hessian",
5946                            fixed_fused_direct.h()[primary][other],
5947                            fixed_fused_scalar.h()[primary][other],
5948                        ),
5949                        (
5950                            "dynamic product Hessian",
5951                            dynamic_product_direct.h_at(primary, other),
5952                            dynamic_product_scalar.h_at(primary, other),
5953                        ),
5954                        (
5955                            "dynamic affine Hessian",
5956                            dynamic_affine_direct.h_at(primary, other),
5957                            dynamic_affine_scalar.h_at(primary, other),
5958                        ),
5959                        (
5960                            "dynamic affine sum Hessian",
5961                            dynamic_sum_direct.h_at(primary, other),
5962                            dynamic_sum_scalar.h_at(primary, other),
5963                        ),
5964                        (
5965                            "dynamic fused Hessian",
5966                            dynamic_fused_direct.h_at(primary, other),
5967                            dynamic_fused_scalar.h_at(primary, other),
5968                        ),
5969                    ] {
5970                        close(actual, expected, case, label);
5971                    }
5972                }
5973            }
5974        }
5975    }
5976
5977    #[test]
5978    fn shared_product_composition_accepts_empty_expression() {
5979        const K: usize = 4;
5980        let fixed_terms: [&Order2<K>; 0] = [];
5981        let fixed_shared = Order2::constant(1.0);
5982        let scales: [f64; 0] = [];
5983        let stacks: [[f64; 5]; 0] = [];
5984        let fixed = Order2::shared_multiply_add_affine_composed_sum(
5985            &fixed_terms,
5986            &fixed_shared,
5987            &fixed_shared,
5988            &scales,
5989            &scales,
5990            &stacks,
5991        );
5992        assert_eq!(fixed.value().to_bits(), 0.0_f64.to_bits());
5993        assert!(fixed.g().iter().all(|&channel| channel == 0.0));
5994        assert!(fixed.h().iter().flatten().all(|&channel| channel == 0.0));
5995
5996        let value_terms: [&RuntimeValue; 0] = [];
5997        let value_shared = RuntimeValue::constant(1.0, K, &());
5998        let value = RuntimeValue::shared_multiply_add_affine_composed_sum(
5999            &value_terms,
6000            &value_shared,
6001            &value_shared,
6002            &scales,
6003            &scales,
6004            &stacks,
6005            K,
6006            &(),
6007        );
6008        assert_eq!(value.value().to_bits(), 0.0_f64.to_bits());
6009        assert_eq!(value.dimension(), K);
6010
6011        let arena = DynamicJetArena::new();
6012        let dynamic_terms: [&DynamicOrder2<'_>; 0] = [];
6013        let dynamic_shared = DynamicOrder2::constant(1.0, K, &arena);
6014        let dynamic = DynamicOrder2::shared_multiply_add_affine_composed_sum(
6015            &dynamic_terms,
6016            &dynamic_shared,
6017            &dynamic_shared,
6018            &scales,
6019            &scales,
6020            &stacks,
6021            K,
6022            &arena,
6023        );
6024        assert_eq!(dynamic.value().to_bits(), 0.0_f64.to_bits());
6025        assert!((0..K).all(|axis| dynamic.g()[axis] == 0.0));
6026        assert!((0..K).all(|row| (0..K).all(|column| dynamic.h_at(row, column) == 0.0)));
6027    }
6028
6029    #[test]
6030    fn runtime_fused_product_composition_preserves_tower4_channels() {
6031        const K: usize = 2;
6032        const N: usize = 9;
6033        let vars = [
6034            Tower4::<K>::variable(0.37, 0),
6035            Tower4::<K>::variable(-0.61, 1),
6036        ];
6037        let upstream = [
6038            vars[0].mul(&vars[1]).add(&vars[0].exp()),
6039            vars[1].mul(&vars[1]).add(&vars[0].scale(0.3)),
6040            vars[0].mul(&vars[0]).sub(&vars[1].scale(-0.2)),
6041        ];
6042        let mut lefts: [Tower4<K>; N] = std::array::from_fn(|term| upstream[term % upstream.len()]);
6043        lefts[4] = lefts[0];
6044        lefts[5] = lefts[0];
6045        let right = upstream[1];
6046        let addend = upstream[2];
6047        let addend_scales: [f64; N] = std::array::from_fn(|term| [-0.0, 1.0, -0.7, 0.25][term % 4]);
6048        let mut input_scales: [f64; N] =
6049            std::array::from_fn(|term| [0.0, -1.3, 0.45, 1.1][term % 4]);
6050        input_scales[0] = -1.1;
6051        input_scales[4] = 1.1;
6052        let stacks: [[f64; 5]; N] = std::array::from_fn(|term| {
6053            let t = term as f64 + 1.0;
6054            [0.17 * t, -0.11 * t, 0.07 * t, -0.03 * t, 0.013 * t]
6055        });
6056
6057        let expected = (0..N).fold(Tower4::<K>::constant(0.0), |sum, term| {
6058            let inner = if addend_scales[term] == 0.0 {
6059                lefts[term].mul(&right)
6060            } else if addend_scales[term] == 1.0 {
6061                JetScalar::multiply_add(&lefts[term], &right, &addend)
6062            } else {
6063                JetScalar::multiply_add(&lefts[term], &right, &addend.scale(addend_scales[term]))
6064            };
6065            sum.add(&JetScalar::affine_compose(
6066                &inner,
6067                input_scales[term],
6068                0.0,
6069                stacks[term],
6070            ))
6071        });
6072        let wrapped_lefts: [FixedRuntimeJet<Tower4<K>, K>; N] =
6073            std::array::from_fn(|term| FixedRuntimeJet::from_inner(lefts[term]));
6074        let wrapped_right = FixedRuntimeJet::from_inner(right);
6075        let wrapped_addend = FixedRuntimeJet::from_inner(addend);
6076        let mut wrapped_left_refs: [&FixedRuntimeJet<Tower4<K>, K>; N] =
6077            std::array::from_fn(|term| &wrapped_lefts[term]);
6078        wrapped_left_refs[4] = &wrapped_lefts[0];
6079        wrapped_left_refs[5] = &wrapped_lefts[0];
6080        let actual = FixedRuntimeJet::<Tower4<K>, K>::shared_multiply_add_affine_composed_sum(
6081            &wrapped_left_refs,
6082            &wrapped_right,
6083            &wrapped_addend,
6084            &addend_scales,
6085            &input_scales,
6086            &stacks,
6087            K,
6088            &(),
6089        )
6090        .into_inner();
6091
6092        let same = |label: &str, got: f64, want: f64| {
6093            let tolerance = 2.0e-13 * got.abs().max(want.abs()).max(1.0);
6094            assert!(
6095                (got - want).abs() <= tolerance,
6096                "{label}: got={got:+.17e}, want={want:+.17e}, tolerance={tolerance:.3e}"
6097            );
6098        };
6099        same("value", actual.v, expected.v);
6100        for a in 0..K {
6101            same("gradient", actual.g[a], expected.g[a]);
6102            for b in 0..K {
6103                same("Hessian", actual.h[a][b], expected.h[a][b]);
6104                for c in 0..K {
6105                    same("third", actual.t3[a][b][c], expected.t3[a][b][c]);
6106                    for d in 0..K {
6107                        same("fourth", actual.t4[a][b][c][d], expected.t4[a][b][c][d]);
6108                    }
6109                }
6110            }
6111        }
6112    }
6113
6114    /// A small polynomial-plus-unary row expression written ONCE, generically
6115    /// over `S: JetScalar<2>`, so it can be evaluated against every scalar:
6116    /// `ℓ = (e^{p0·p1} + 2) · √(p0·p0 + 1) − p1·p1·0.5`.
6117    /// Exercises mul, add/sub, scale, exp, sqrt — every algebra op.
6118    fn row_expr<S: JetScalar<2>>(p: &[S; 2]) -> S {
6119        let g = p[0].mul(&p[1]).exp();
6120        let inner = g.add(&S::constant(2.0));
6121        let radic = p[0].mul(&p[0]).add(&S::constant(1.0)).sqrt();
6122        inner.mul(&radic).sub(&p[1].mul(&p[1]).scale(0.5))
6123    }
6124
6125    /// The same expression exposed through the canonical generic row-program seam.
6126    struct ExprProgram {
6127        p: [f64; 2],
6128    }
6129    impl RowProgram<2> for ExprProgram {
6130        fn n_rows(&self) -> usize {
6131            1
6132        }
6133        fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
6134            if row >= self.n_rows() {
6135                return Err(format!("ExprProgram: row {row} out of range"));
6136            }
6137            Ok(self.p)
6138        }
6139        fn eval<S: JetScalar<2>>(&self, row: usize, p: &[S; 2]) -> Result<S, String> {
6140            if row >= self.n_rows() {
6141                return Err(format!("ExprProgram: row {row} out of range"));
6142            }
6143            Ok(row_expr(p))
6144        }
6145    }
6146
6147    const SEED: [f64; 2] = [0.37, -0.81];
6148    const U: [f64; 2] = [0.6, -0.2];
6149    const V: [f64; 2] = [-0.4, 1.1];
6150    const TOL: f64 = 1e-10;
6151
6152    fn close(a: f64, b: f64, label: &str) {
6153        let band = TOL + TOL * a.abs().max(b.abs());
6154        assert!(
6155            (a - b).abs() <= band,
6156            "{label}: {a:+.15e} vs {b:+.15e} (band {band:.3e})"
6157        );
6158    }
6159
6160    fn tower() -> Tower4<2> {
6161        *program_full_tower(&ExprProgram { p: SEED }, 0).expect("tower")
6162    }
6163
6164    /// Order2 reproduces Tower4's value/grad/Hessian channels exactly.
6165    #[test]
6166    fn order2_matches_tower_value_grad_hessian() {
6167        let t = tower();
6168        let vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6169        let s = row_expr(&vars);
6170        close(s.value(), t.v, "value");
6171        for a in 0..2 {
6172            close(s.0.g[a], t.g[a], &format!("grad[{a}]"));
6173            for b in 0..2 {
6174                close(s.h()[a][b], t.h[a][b], &format!("hess[{a}][{b}]"));
6175            }
6176        }
6177    }
6178
6179    #[test]
6180    fn mapped_order2_accumulator_matches_dense_overlapping_atoms() {
6181        const K: usize = 4;
6182        let p = [0.2_f64, 0.7, -0.4, 0.3];
6183        let dense_vars: [Order2<K>; K] =
6184            std::array::from_fn(|axis| Order2::variable(p[axis], axis));
6185        let dense_q0 = dense_vars[3].mul(&dense_vars[1]).add(&dense_vars[3].exp());
6186        let dense_q1 = dense_vars[1].mul(&dense_vars[2]).sub(&dense_vars[2]);
6187        let dense = dense_q0.ln().add(&dense_q1.exp());
6188
6189        let local_q0_vars: [Order2<2>; 2] =
6190            std::array::from_fn(|axis| Order2::variable(p[[3, 1][axis]], axis));
6191        let local_q0 = local_q0_vars[0]
6192            .mul(&local_q0_vars[1])
6193            .add(&local_q0_vars[0].exp());
6194        let local_q1_vars: [Order2<2>; 2] =
6195            std::array::from_fn(|axis| Order2::variable(p[[1, 2][axis]], axis));
6196        let local_q1 = local_q1_vars[0]
6197            .mul(&local_q1_vars[1])
6198            .sub(&local_q1_vars[1]);
6199
6200        let q0 = local_q0.value();
6201        let q1_exp = local_q1.value().exp();
6202        let mut lowered = MappedOrder2Accumulator::<K>::zero();
6203        lowered.add_composed(
6204            &local_q0,
6205            [3, 1],
6206            [q0.ln(), q0.recip(), -1.0 / (q0 * q0)],
6207            false,
6208            [false, false],
6209            [false, false, false],
6210        );
6211        lowered.add_composed(
6212            &local_q1,
6213            [1, 2],
6214            [q1_exp, q1_exp, q1_exp],
6215            true,
6216            [true, false],
6217            [true, false, false],
6218        );
6219        let (value, gradient, hessian) = lowered.into_channels();
6220
6221        close(value, dense.value(), "mapped value");
6222        for i in 0..K {
6223            close(gradient[i], dense.g()[i], &format!("mapped gradient[{i}]"));
6224            for j in 0..K {
6225                close(
6226                    hessian[i][j],
6227                    dense.h()[i][j],
6228                    &format!("mapped Hessian[{i},{j}]"),
6229                );
6230            }
6231        }
6232    }
6233
6234    #[test]
6235    #[should_panic(expected = "mapped atom axes must be injective")]
6236    fn mapped_order2_accumulator_rejects_duplicate_axes() {
6237        let vars: [Order2<2>; 2] = std::array::from_fn(|axis| Order2::variable(0.2, axis));
6238        let atom = vars[0].add(&vars[1]);
6239        let mut lowered = MappedOrder2Accumulator::<2>::zero();
6240        lowered.add_composed(
6241            &atom,
6242            [1, 1],
6243            [0.4, 1.0, 0.0],
6244            false,
6245            [false, false],
6246            [false, false, false],
6247        );
6248    }
6249
6250    #[test]
6251    #[should_panic(expected = "mapped atom axis must be within")]
6252    fn mapped_order2_accumulator_rejects_out_of_range_axes() {
6253        let atom = Order2::<1>::variable(0.2, 0);
6254        let mut lowered = MappedOrder2Accumulator::<2>::zero();
6255        lowered.add_composed(&atom, [2], [0.2, 1.0, 0.0], false, [false], [false]);
6256    }
6257
6258    #[test]
6259    fn dynamic_order2_accumulator_matches_dense_composed_sum() {
6260        const K: usize = 4;
6261
6262        struct Term {
6263            first: f64,
6264            second: f64,
6265            gradient: [f64; K],
6266            hessian: [[f64; K]; K],
6267        }
6268
6269        impl DynamicOrder2Term for Term {
6270            fn outer_first(&self) -> f64 {
6271                self.first
6272            }
6273
6274            fn outer_second(&self) -> f64 {
6275                self.second
6276            }
6277
6278            fn inner_gradient(&self, axis: usize) -> f64 {
6279                self.gradient[axis]
6280            }
6281
6282            fn inner_hessian(&self, row: usize, column: usize) -> f64 {
6283                self.hessian[row][column]
6284            }
6285        }
6286
6287        let p = [0.7, -0.3, 0.2, 0.8];
6288        let vars: [Order2<K>; K] = std::array::from_fn(|axis| Order2::variable(p[axis], axis));
6289        let first_atom = vars[0]
6290            .mul(&vars[1])
6291            .add(&vars[2].exp())
6292            .add(&Order2::constant(1.5));
6293        let second_atom = vars[1].mul(&vars[3]).sub(&vars[0]);
6294        let first_value = first_atom.value();
6295        let second_exp = second_atom.value().exp();
6296        let first_stack = [
6297            first_value.ln(),
6298            first_value.recip(),
6299            -1.0 / (first_value * first_value),
6300            0.0,
6301            0.0,
6302        ];
6303        let second_stack = [second_exp, second_exp, second_exp, second_exp, second_exp];
6304        let dense = first_atom
6305            .compose_unary(first_stack)
6306            .add(&second_atom.compose_unary(second_stack));
6307        let terms = [
6308            Term {
6309                first: first_stack[1],
6310                second: first_stack[2],
6311                gradient: *first_atom.g(),
6312                hessian: *first_atom.h(),
6313            },
6314            Term {
6315                first: second_stack[1],
6316                second: second_stack[2],
6317                gradient: *second_atom.g(),
6318                hessian: *second_atom.h(),
6319            },
6320        ];
6321        let (value, gradient, hessian) = DynamicOrder2Accumulator::from_composed_sum(
6322            K,
6323            first_stack[0] + second_stack[0],
6324            &terms,
6325        )
6326        .into_channels();
6327
6328        close(value, dense.value(), "dynamic value");
6329        for row in 0..K {
6330            close(
6331                gradient[row],
6332                dense.g()[row],
6333                &format!("dynamic gradient[{row}]"),
6334            );
6335            for column in 0..K {
6336                close(
6337                    hessian[row * K + column],
6338                    dense.h()[row][column],
6339                    &format!("dynamic Hessian[{row},{column}]"),
6340                );
6341            }
6342        }
6343    }
6344
6345    #[derive(Clone, Copy, Debug)]
6346    struct FullTwoPattern;
6347
6348    impl HessianPattern<2, 3> for FullTwoPattern {
6349        const PAIRS: [(usize, usize); 3] = [(0, 0), (0, 1), (1, 1)];
6350        const PAIR_BITS: [[u128; 2]; 2] = hessian_pair_bits(Self::PAIRS);
6351    }
6352
6353    /// Patterned order-two arithmetic is channel-identical to dense `Order2`
6354    /// when the pattern covers the expression's complete Hessian support.
6355    #[test]
6356    fn patterned_order2_matches_dense_order2() {
6357        type Sparse = PatternedOrder2<FullTwoPattern, 2, 3>;
6358        let dense_vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6359        let sparse_vars: [Sparse; 2] = std::array::from_fn(|a| Sparse::variable(SEED[a], a));
6360        let dense = row_expr(&dense_vars);
6361        let sparse = row_expr(&sparse_vars);
6362        close(sparse.value(), dense.value(), "patterned value");
6363        for i in 0..2 {
6364            close(sparse.g()[i], dense.g()[i], &format!("patterned grad[{i}]"));
6365            for j in 0..2 {
6366                close(
6367                    sparse.h()[i][j],
6368                    dense.h()[i][j],
6369                    &format!("patterned hess[{i}][{j}]"),
6370                );
6371            }
6372        }
6373    }
6374
6375    /// `compose_unary_with` is `to_bits`-identical to the explicit
6376    /// `compose_unary(stack_fn(value))` form. Exercised on [`Order2`] across
6377    /// `K ∈ {2,3,4,9}`, ≥ 4000 random seeded inputs.
6378    #[test]
6379    fn compose_unary_with_scalar_seam_bit_identical() {
6380        fn rand_unit(state: &mut u64) -> f64 {
6381            let mut x = *state;
6382            x ^= x << 13;
6383            x ^= x >> 7;
6384            x ^= x << 17;
6385            *state = x;
6386            2.0 * ((x >> 11) as f64 / ((1u64 << 53) as f64)) - 1.0
6387        }
6388        // A base-value-dependent finite stack standing in for a family stack.
6389        fn stack(u: f64) -> [f64; 5] {
6390            [
6391                u.sin(),
6392                u.cos(),
6393                (2.0 * u).sin(),
6394                (0.5 * u).cos(),
6395                u * u - 0.3,
6396            ]
6397        }
6398        fn run<const K: usize>(state: &mut u64, n: usize) -> usize {
6399            for _ in 0..n {
6400                // A non-trivial Order2<K> jet: a seeded variable pushed through a
6401                // couple of algebra ops so g/h are dense, then exercise the seam.
6402                let base = rand_unit(state);
6403                let mut s = Order2::<K>::variable(base, 0);
6404                for a in 1..K {
6405                    s = crate::nested_dual::JetField::mul(
6406                        &s,
6407                        &Order2::<K>::variable(rand_unit(state), a),
6408                    );
6409                }
6410                let with = s.compose_unary_with(stack);
6411                let explicit = s.compose_unary(stack(s.value()));
6412                assert_eq!(with.value().to_bits(), explicit.value().to_bits(), "value");
6413                for a in 0..K {
6414                    assert_eq!(with.g()[a].to_bits(), explicit.g()[a].to_bits(), "g[{a}]");
6415                    for b in 0..K {
6416                        assert_eq!(
6417                            with.h()[a][b].to_bits(),
6418                            explicit.h()[a][b].to_bits(),
6419                            "h[{a}][{b}]"
6420                        );
6421                    }
6422                }
6423            }
6424            n
6425        }
6426        let mut st = 0x9e37_79b9_7f4a_7c15u64;
6427        let total = run::<2>(&mut st, 1100)
6428            + run::<3>(&mut st, 1100)
6429            + run::<4>(&mut st, 1100)
6430            + run::<9>(&mut st, 1100);
6431        assert_eq!(total, 4400);
6432    }
6433
6434    /// OneSeed's ε-Hessian is the contracted third Σ_c ℓ_{abc} u_c, matching
6435    /// `Tower4::third_contracted(u)`. Base channels also match the tower.
6436    #[test]
6437    fn one_seed_matches_tower_third_contracted() {
6438        let t = tower();
6439        let truth = t.third_contracted(&U);
6440        let vars: [OneSeed<2>; 2] =
6441            std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
6442        let s = row_expr(&vars);
6443        // Base channels are the plain (v, g, H).
6444        close(s.value(), t.v, "value");
6445        for a in 0..2 {
6446            for b in 0..2 {
6447                close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
6448            }
6449        }
6450        let third = s.contracted_third();
6451        for a in 0..2 {
6452            for b in 0..2 {
6453                close(third[a][b], truth[a][b], &format!("third[{a}][{b}]"));
6454            }
6455        }
6456    }
6457
6458    /// The fused fixed-width OneSeed product/composition preserves every
6459    /// channel of the defining unfused nilpotent algebra. Random symmetric
6460    /// Order2 inputs exercise dense non-seed intermediates, while exact
6461    /// mirrored symmetry is checked as an independent production invariant.
6462    #[test]
6463    fn fused_one_seed_channels_match_unfused_definition_932() {
6464        const K: usize = 8;
6465
6466        fn random_scalar(state: &mut u64) -> f64 {
6467            *state ^= *state << 13;
6468            *state ^= *state >> 7;
6469            *state ^= *state << 17;
6470            ((*state >> 11) as f64 / ((1_u64 << 53) as f64)) * 2.0 - 1.0
6471        }
6472
6473        fn random_order2<const N: usize>(state: &mut u64) -> Order2<N> {
6474            let mut tower = crate::jet_tower::Tower2::<N>::zero();
6475            tower.v = random_scalar(state);
6476            for axis in 0..N {
6477                tower.g[axis] = random_scalar(state);
6478            }
6479            for row in 0..N {
6480                for column in row..N {
6481                    let channel = random_scalar(state);
6482                    tower.h[row][column] = channel;
6483                    tower.h[column][row] = channel;
6484                }
6485            }
6486            Order2(tower)
6487        }
6488
6489        fn assert_channels_close<const N: usize>(
6490            label: &str,
6491            actual: &OneSeed<N>,
6492            expected: &OneSeed<N>,
6493        ) {
6494            for (part_label, actual_part, expected_part, require_exact_symmetry) in [
6495                ("base", &actual.base.0, &expected.base.0, false),
6496                ("eps", &actual.eps.0, &expected.eps.0, true),
6497            ] {
6498                let check = |channel: &str, got: f64, want: f64| {
6499                    let tolerance = 2.0e-14 * got.abs().max(want.abs()).max(1.0);
6500                    assert!(
6501                        (got - want).abs() <= tolerance,
6502                        "{label} {part_label} {channel}: got={got:+.17e} want={want:+.17e}"
6503                    );
6504                };
6505                check("value", actual_part.v, expected_part.v);
6506                for row in 0..N {
6507                    check(
6508                        &format!("gradient[{row}]"),
6509                        actual_part.g[row],
6510                        expected_part.g[row],
6511                    );
6512                    for column in 0..N {
6513                        check(
6514                            &format!("hessian[{row},{column}]"),
6515                            actual_part.h[row][column],
6516                            expected_part.h[row][column],
6517                        );
6518                        if require_exact_symmetry {
6519                            assert_eq!(
6520                                actual_part.h[row][column].to_bits(),
6521                                actual_part.h[column][row].to_bits(),
6522                                "{label} {part_label} Hessian symmetry at [{row},{column}]"
6523                            );
6524                        }
6525                    }
6526                }
6527            }
6528        }
6529
6530        let mut state = 0x9320_1eed_5eed_cafe_u64;
6531        for sample in 0..256 {
6532            let left = OneSeed {
6533                base: random_order2::<K>(&mut state),
6534                eps: random_order2::<K>(&mut state),
6535            };
6536            let right = OneSeed {
6537                base: random_order2::<K>(&mut state),
6538                eps: random_order2::<K>(&mut state),
6539            };
6540
6541            let fused_product = left.mul(&right);
6542            let unfused_product = OneSeed {
6543                base: left.base.mul(&right.base),
6544                eps: left.base.mul(&right.eps).add(&left.eps.mul(&right.base)),
6545            };
6546            assert_channels_close(
6547                &format!("sample {sample} product"),
6548                &fused_product,
6549                &unfused_product,
6550            );
6551
6552            let derivatives: [f64; 5] = std::array::from_fn(|_| random_scalar(&mut state));
6553            let fused_composition = left.compose_unary(derivatives);
6554            let unfused_composition = OneSeed {
6555                base: left.base.compose_unary(derivatives),
6556                eps: left
6557                    .base
6558                    .compose_unary([
6559                        derivatives[1],
6560                        derivatives[2],
6561                        derivatives[3],
6562                        derivatives[4],
6563                        derivatives[4],
6564                    ])
6565                    .mul(&left.eps),
6566            };
6567            assert_channels_close(
6568                &format!("sample {sample} composition"),
6569                &fused_composition,
6570                &unfused_composition,
6571            );
6572        }
6573    }
6574
6575    /// TwoSeed's εδ-Hessian is the contracted fourth Σ_{cd} ℓ_{abcd} u_c v_d,
6576    /// matching `Tower4::fourth_contracted(u, v)`. The ε / δ single-seed parts
6577    /// reproduce the two third contractions Σ_c ℓ_{abc} u_c and …v_d.
6578    #[test]
6579    fn two_seed_matches_tower_fourth_contracted() {
6580        let t = tower();
6581        let truth4 = t.fourth_contracted(&U, &V);
6582        let truth3_u = t.third_contracted(&U);
6583        let truth3_v = t.third_contracted(&V);
6584        let vars: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
6585        let s = row_expr(&vars);
6586        close(s.value(), t.v, "value");
6587        for a in 0..2 {
6588            close(s.base.0.g[a], t.g[a], &format!("grad[{a}]"));
6589            for b in 0..2 {
6590                close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
6591                close(
6592                    s.eps.h()[a][b],
6593                    truth3_u[a][b],
6594                    &format!("eps third_u[{a}][{b}]"),
6595                );
6596                close(
6597                    s.del.h()[a][b],
6598                    truth3_v[a][b],
6599                    &format!("del third_v[{a}][{b}]"),
6600                );
6601            }
6602        }
6603        let fourth = s.contracted_fourth();
6604        for a in 0..2 {
6605            for b in 0..2 {
6606                close(fourth[a][b], truth4[a][b], &format!("fourth[{a}][{b}]"));
6607            }
6608        }
6609    }
6610
6611    /// The generic `row_nll_generic` seam (added to Tower4's program trait
6612    /// surface) evaluates the SAME expression on each scalar and extracts the
6613    /// channel a consumer asks for, agreeing with the direct Tower4 contraction.
6614    #[test]
6615    fn generic_program_seam_matches_tower_for_every_channel() {
6616        let t = tower();
6617        // Order2 via generic seam.
6618        let o2: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6619        let so2 = row_expr(&o2);
6620        close(so2.value(), t.v, "seam order2 value");
6621        // OneSeed third.
6622        let os: [OneSeed<2>; 2] =
6623            std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
6624        let third = row_expr(&os).contracted_third();
6625        let truth3 = t.third_contracted(&U);
6626        for a in 0..2 {
6627            for b in 0..2 {
6628                close(third[a][b], truth3[a][b], &format!("seam third[{a}][{b}]"));
6629            }
6630        }
6631        // TwoSeed fourth.
6632        let ts: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
6633        let fourth = row_expr(&ts).contracted_fourth();
6634        let truth4 = t.fourth_contracted(&U, &V);
6635        for a in 0..2 {
6636            for b in 0..2 {
6637                close(
6638                    fourth[a][b],
6639                    truth4[a][b],
6640                    &format!("seam fourth[{a}][{b}]"),
6641                );
6642            }
6643        }
6644    }
6645
6646    /// The (test-only) `Tower4: JetScalar` impl is the all-channels oracle scalar:
6647    /// evaluating the SAME generic `row_expr` at `S = Tower4` (through the
6648    /// `JetScalar` trait ops) must reproduce, channel-for-channel, the `Tower4`
6649    /// obtained through the canonical `RowProgram` path (`program_full_tower`).
6650    /// This pins that the full-tower scalar and the contracted scalar oracles above
6651    /// all execute the same generic algebra surface.
6652    #[test]
6653    fn tower4_as_jetscalar_matches_program_tower_all_channels() {
6654        let t = tower();
6655        let vars: [Tower4<2>; 2] = std::array::from_fn(|a| Tower4::variable(SEED[a], a));
6656        let s = row_expr(&vars);
6657        close(s.v, t.v, "tower-jetscalar value");
6658        for a in 0..2 {
6659            close(s.g[a], t.g[a], &format!("tower-jetscalar grad[{a}]"));
6660            for b in 0..2 {
6661                close(
6662                    s.h[a][b],
6663                    t.h[a][b],
6664                    &format!("tower-jetscalar hess[{a}][{b}]"),
6665                );
6666                for c in 0..2 {
6667                    close(
6668                        s.t3[a][b][c],
6669                        t.t3[a][b][c],
6670                        &format!("tower-jetscalar t3[{a}][{b}][{c}]"),
6671                    );
6672                    for d in 0..2 {
6673                        close(
6674                            s.t4[a][b][c][d],
6675                            t.t4[a][b][c][d],
6676                            &format!("tower-jetscalar t4[{a}][{b}][{c}][{d}]"),
6677                        );
6678                    }
6679                }
6680            }
6681        }
6682    }
6683
6684    /// Arena-backed runtime directional jets must reproduce the fixed packed
6685    /// algebras channel-for-channel. This isolates the runtime scalar algebra
6686    /// from every family row program before the SLS/SAE integration oracles.
6687    #[test]
6688    fn runtime_directional_jets_match_fixed_packed_algebra_932() {
6689        fn expression<'arena, S: RuntimeJetScalar<'arena>>(vars: &[S]) -> S {
6690            let bilinear = vars[0].mul(&vars[1]);
6691            let curved = vars[2].scale(0.7).add(&vars[3].mul(&vars[3]).scale(-0.2));
6692            bilinear
6693                .add(&curved)
6694                .exp()
6695                .mul(&vars[4].compose_unary([0.4, -0.3, 0.2, -0.1, 0.05]))
6696        }
6697
6698        const K: usize = 5;
6699        let values = [0.2, -0.7, 0.4, 1.1, -0.3];
6700        let direction_u = [0.5, -0.2, 0.7, -0.4, 0.1];
6701        let direction_v = [-0.3, 0.8, 0.2, 0.6, -0.5];
6702        let close = |actual: f64, expected: f64| {
6703            let tolerance = 1.0e-13 * (1.0 + actual.abs().max(expected.abs()));
6704            assert!((actual - expected).abs() <= tolerance);
6705        };
6706
6707        let fixed_one: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
6708            .map(|axis| FixedRuntimeJet {
6709                inner: OneSeed::seed_direction(values[axis], axis, direction_u[axis]),
6710            })
6711            .collect();
6712        let arena_one = DynamicJetArena::new();
6713        let dynamic_one: Vec<DynamicOneSeed<'_>> = (0..K)
6714            .map(|axis| {
6715                DynamicOneSeed::seed_direction(values[axis], axis, direction_u[axis], K, &arena_one)
6716            })
6717            .collect();
6718        let fixed_third = expression(&fixed_one).into_inner().contracted_third();
6719        let dynamic_third = expression(&dynamic_one);
6720        for a in 0..K {
6721            for b in 0..K {
6722                assert_eq!(
6723                    dynamic_third.contracted_third()[a * K + b].to_bits(),
6724                    dynamic_third.contracted_third()[b * K + a].to_bits(),
6725                    "arena third Hessian must be exactly symmetric at ({a},{b})"
6726                );
6727                close(
6728                    dynamic_third.contracted_third()[a * K + b],
6729                    fixed_third[a][b],
6730                );
6731            }
6732        }
6733
6734        let fixed_one_v: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
6735            .map(|axis| FixedRuntimeJet {
6736                inner: OneSeed::seed_direction(values[axis], axis, direction_v[axis]),
6737            })
6738            .collect();
6739        let fixed_third_v = expression(&fixed_one_v).into_inner().contracted_third();
6740        let batch_workspace = DynamicJetBatchWorkspace::new(2);
6741        let directions = [direction_u, direction_v];
6742        let batch_vars = batch_workspace.alloc_slice_fill_with(K, |axis| {
6743            DynamicOneSeedBatch::seed_directions(values[axis], axis, K, &batch_workspace, |lane| {
6744                directions[lane][axis]
6745            })
6746        });
6747        let dynamic_batch = expression(batch_vars);
6748        assert_eq!(dynamic_batch.lanes(), 2);
6749        for lane in 0..2 {
6750            let expected = if lane == 0 {
6751                &fixed_third
6752            } else {
6753                &fixed_third_v
6754            };
6755            for a in 0..K {
6756                for b in 0..K {
6757                    close(
6758                        dynamic_batch.contracted_third(lane)[a * K + b],
6759                        expected[a][b],
6760                    );
6761                }
6762            }
6763        }
6764
6765        let fixed_two: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
6766            .map(|axis| FixedRuntimeJet {
6767                inner: TwoSeed::seed(values[axis], axis, direction_u[axis], direction_v[axis]),
6768            })
6769            .collect();
6770        let arena_two = DynamicJetArena::new();
6771        let dynamic_two: Vec<DynamicTwoSeed<'_>> = (0..K)
6772            .map(|axis| {
6773                DynamicTwoSeed::seed(
6774                    values[axis],
6775                    axis,
6776                    direction_u[axis],
6777                    direction_v[axis],
6778                    K,
6779                    &arena_two,
6780                )
6781            })
6782            .collect();
6783        let fixed_fourth = expression(&fixed_two).into_inner().contracted_fourth();
6784        let dynamic_fourth = expression(&dynamic_two);
6785        for a in 0..K {
6786            for b in 0..K {
6787                close(
6788                    dynamic_fourth.contracted_fourth()[a * K + b],
6789                    fixed_fourth[a][b],
6790                );
6791            }
6792        }
6793
6794        let fixed_two_swapped: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
6795            .map(|axis| {
6796                FixedRuntimeJet::from_inner(TwoSeed::seed(
6797                    values[axis],
6798                    axis,
6799                    direction_v[axis],
6800                    direction_u[axis],
6801                ))
6802            })
6803            .collect();
6804        let fixed_fourth_swapped = expression(&fixed_two_swapped)
6805            .into_inner()
6806            .contracted_fourth();
6807        let pair_workspace = DynamicJetBatchWorkspace::new(2);
6808        let direction_pairs = [(direction_u, direction_v), (direction_v, direction_u)];
6809        let pair_vars = pair_workspace.alloc_slice_fill_with(K, |axis| {
6810            DynamicTwoSeedBatch::seed_direction_pairs(
6811                values[axis],
6812                axis,
6813                K,
6814                &pair_workspace,
6815                |lane| (direction_pairs[lane].0[axis], direction_pairs[lane].1[axis]),
6816            )
6817        });
6818        let dynamic_pair_batch = expression(pair_vars);
6819        assert_eq!(dynamic_pair_batch.lanes(), 2);
6820        for lane in 0..2 {
6821            let expected = if lane == 0 {
6822                &fixed_fourth
6823            } else {
6824                &fixed_fourth_swapped
6825            };
6826            for a in 0..K {
6827                for b in 0..K {
6828                    close(
6829                        dynamic_pair_batch.contracted_fourth(lane)[a * K + b],
6830                        expected[a][b],
6831                    );
6832                }
6833            }
6834        }
6835    }
6836
6837    #[test]
6838    fn dynamic_jet_arena_compacts_fragmented_high_water_932() {
6839        const WORDS_PER_ALLOCATION: usize = 1 << 17;
6840        const ALLOCATIONS: usize = 6;
6841
6842        let mut arena = DynamicJetArena::new();
6843        for lane in 0..ALLOCATIONS {
6844            let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
6845            std::hint::black_box(allocation);
6846        }
6847        let fragmented_high_water = arena.allocated_bytes();
6848
6849        arena.reset();
6850        let compact_high_water = arena.allocated_bytes();
6851        assert!(
6852            compact_high_water >= fragmented_high_water,
6853            "compacted arena must retain the complete fragmented tape"
6854        );
6855
6856        for lane in 0..ALLOCATIONS {
6857            let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
6858            std::hint::black_box(allocation);
6859        }
6860        assert_eq!(
6861            arena.allocated_bytes(),
6862            compact_high_water,
6863            "equal replay must fit in the compacted chunk"
6864        );
6865
6866        arena.reset();
6867        assert_eq!(
6868            arena.allocated_bytes(),
6869            compact_high_water,
6870            "stable reset must retain the compacted chunk"
6871        );
6872    }
6873}
6874
6875#[cfg(test)]
6876mod batch_tests {
6877    //! SIMD row-batching oracle: prove [`Order2Batch<K>`] (4 rows in
6878    //! `wide::f64x4` lanes) is `to_bits`-identical, on every value/gradient/
6879    //! Hessian channel, to the production [`Order2<K>`] evaluated per row — and
6880    //! that the new scalar field [`Order2Lane<f64, K>`] is too. Composing the two
6881    //! claims, batch lane `i` reproduces the production scalar for row `i` bit
6882    //! for bit, so the 4× throughput is a free lunch (no result change).
6883
6884    use super::{
6885        JetScalar, Lane, OneSeed, OneSeedBatch, OneSeedLane, Order2, Order2Batch, Order2Lane,
6886        TwoSeed, TwoSeedBatch, TwoSeedLane,
6887    };
6888    // The scalar-field algebra (`value`, `add`, …) lives on the shared `JetField`
6889    // base now, so the concrete-typed channel reads below need it in scope.
6890    use crate::nested_dual::JetField;
6891
6892    /// The ops the witness row expression needs, so ONE generic body evaluates
6893    /// at the production [`Order2<K>`], the new scalar [`Order2Lane<f64, K>`],
6894    /// and the batched [`Order2Batch<K>`].
6895    trait RowAlg<const K: usize>: Copy {
6896        fn constant(c: f64) -> Self;
6897        fn add(&self, o: &Self) -> Self;
6898        fn sub(&self, o: &Self) -> Self;
6899        fn mul(&self, o: &Self) -> Self;
6900        fn scale(&self, s: f64) -> Self;
6901        fn exp(&self) -> Self;
6902        fn sqrt(&self) -> Self;
6903        fn recip(&self) -> Self;
6904    }
6905
6906    impl<const K: usize> RowAlg<K> for Order2<K> {
6907        fn constant(c: f64) -> Self {
6908            <Self as JetScalar<K>>::constant(c)
6909        }
6910        fn add(&self, o: &Self) -> Self {
6911            crate::nested_dual::JetField::add(self, o)
6912        }
6913        fn sub(&self, o: &Self) -> Self {
6914            crate::nested_dual::JetField::sub(self, o)
6915        }
6916        fn mul(&self, o: &Self) -> Self {
6917            crate::nested_dual::JetField::mul(self, o)
6918        }
6919        fn scale(&self, s: f64) -> Self {
6920            crate::nested_dual::JetField::scale(self, s)
6921        }
6922        fn exp(&self) -> Self {
6923            JetScalar::exp(self)
6924        }
6925        fn sqrt(&self) -> Self {
6926            JetScalar::sqrt(self)
6927        }
6928        fn recip(&self) -> Self {
6929            JetScalar::recip(self)
6930        }
6931    }
6932
6933    impl<L: Lane, const K: usize> RowAlg<K> for Order2Lane<L, K> {
6934        fn constant(c: f64) -> Self {
6935            Order2Lane::constant(L::splat(c))
6936        }
6937        fn add(&self, o: &Self) -> Self {
6938            Order2Lane::add(self, o)
6939        }
6940        fn sub(&self, o: &Self) -> Self {
6941            Order2Lane::sub(self, o)
6942        }
6943        fn mul(&self, o: &Self) -> Self {
6944            Order2Lane::mul(self, o)
6945        }
6946        fn scale(&self, s: f64) -> Self {
6947            Order2Lane::scale(self, s)
6948        }
6949        fn exp(&self) -> Self {
6950            Order2Lane::exp(self)
6951        }
6952        fn sqrt(&self) -> Self {
6953            Order2Lane::sqrt(self)
6954        }
6955        fn recip(&self) -> Self {
6956            Order2Lane::recip(self)
6957        }
6958    }
6959
6960    /// A dense witness row expression touching every algebra op (mul, add, sub,
6961    /// scale, exp, sqrt, recip) over ALL `K` primaries, so the gradient and the
6962    /// full `K×K` Hessian are dense (no trivially-zero channel). All transcend.
6963    /// arguments are kept finite/positive: `sqrt(s²+1) > 0`, `recip(exp+2) > 0`.
6964    fn row_expr<const K: usize, A: RowAlg<K>>(p: &[A; K]) -> A {
6965        let mut s = A::constant(0.3);
6966        for a in 0..K {
6967            let b = (a + 1) % K;
6968            s = s.add(&p[a].mul(&p[b]).scale(0.1 + 0.05 * a as f64));
6969        }
6970        let e = s.exp();
6971        let r = s.mul(&s).add(&A::constant(1.0)).sqrt();
6972        let denom = e.add(&A::constant(2.0));
6973        e.mul(&r).sub(&s.scale(0.5)).mul(&denom.recip())
6974    }
6975
6976    /// xorshift64 → `f64` in `[-1, 1)`.
6977    fn rand_unit(state: &mut u64) -> f64 {
6978        let mut x = *state;
6979        x ^= x << 13;
6980        x ^= x >> 7;
6981        x ^= x << 17;
6982        *state = x;
6983        let u = (x >> 11) as f64 / ((1u64 << 53) as f64); // [0, 1)
6984        2.0 * u - 1.0
6985    }
6986
6987    /// Returns the number of (batch, row) pairs whose every channel was
6988    /// verified bit-identical, so the caller can assert the expected total ran.
6989    fn check_k<const K: usize>(state: &mut u64, batches: usize) -> usize {
6990        let mut verified_rows = 0usize;
6991        for _ in 0..batches {
6992            // Four independent rows of K primary values.
6993            let rows: [[f64; K]; 4] =
6994                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
6995
6996            // Production ground truth, evaluated per row at Order2<K>.
6997            let prod: [Order2<K>; 4] = std::array::from_fn(|r| {
6998                let p: [Order2<K>; K] = std::array::from_fn(|a| Order2::variable(rows[r][a], a));
6999                row_expr(&p)
7000            });
7001
7002            // New scalar field (Order2Lane<f64>), per row.
7003            let scal: [Order2Lane<f64, K>; 4] = std::array::from_fn(|r| {
7004                let p: [Order2Lane<f64, K>; K] =
7005                    std::array::from_fn(|a| Order2Lane::variable(rows[r][a], a));
7006                row_expr(&p)
7007            });
7008
7009            // Batched: 4 rows packed into f64x4 lanes, ONE vector pass.
7010            let pbatch: [Order2Batch<K>; K] = std::array::from_fn(|a| {
7011                let packed = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7012                Order2Batch::variable(packed, a)
7013            });
7014            let batch = row_expr(&pbatch);
7015
7016            for r in 0..4 {
7017                let g = prod[r].0;
7018                // Order2Lane<f64> == Order2<K> (bit-identical scalar field).
7019                assert_eq!(scal[r].v.to_bits(), g.v.to_bits(), "K={K} scalar v");
7020                // Batch lane r == Order2<K> for row r.
7021                let lr = batch.lane(r).0;
7022                assert_eq!(lr.v.to_bits(), g.v.to_bits(), "K={K} batch lane {r} v");
7023                for a in 0..K {
7024                    assert_eq!(
7025                        scal[r].g[a].to_bits(),
7026                        g.g[a].to_bits(),
7027                        "K={K} scalar g[{a}]"
7028                    );
7029                    assert_eq!(
7030                        lr.g[a].to_bits(),
7031                        g.g[a].to_bits(),
7032                        "K={K} batch lane {r} g[{a}]"
7033                    );
7034                    for b in 0..K {
7035                        assert_eq!(
7036                            scal[r].h[a][b].to_bits(),
7037                            g.h[a][b].to_bits(),
7038                            "K={K} scalar h[{a}][{b}]"
7039                        );
7040                        assert_eq!(
7041                            lr.h[a][b].to_bits(),
7042                            g.h[a][b].to_bits(),
7043                            "K={K} batch lane {r} h[{a}][{b}]"
7044                        );
7045                    }
7046                }
7047                verified_rows += 1;
7048            }
7049        }
7050        verified_rows
7051    }
7052
7053    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: every channel of
7054    /// every lane is `to_bits`-identical to the production scalar per row.
7055    #[test]
7056    fn batch_lanes_bit_identical_to_scalar_per_row() {
7057        let mut state = 0x9E37_79B9_7F4A_7C15_u64;
7058        let mut verified = 0usize;
7059        verified += check_k::<2>(&mut state, 2000);
7060        verified += check_k::<3>(&mut state, 2000);
7061        verified += check_k::<4>(&mut state, 2000);
7062        verified += check_k::<9>(&mut state, 2000);
7063        // 4 K-values × 2000 batches × 4 packed rows each, all bit-identical.
7064        assert_eq!(verified, 4 * 2000 * 4, "every batch row must be verified");
7065    }
7066
7067    // ── One-/two-seed lane oracles ──────────────────────────────────────────
7068    //
7069    // The same dense `row_expr` witness program runs over the SEEDED directional
7070    // scalars: the scalar `OneSeed`/`TwoSeed` per row, the `f64`-lane re-type
7071    // (`*SeedLane<f64>`), and the 4-rows-per-pass batch (`*SeedBatch`). The
7072    // headline claim is that the contracted-third / contracted-fourth channel of
7073    // every lane is `to_bits`-identical to the production scalar's per row.
7074
7075    impl<const K: usize> RowAlg<K> for OneSeed<K> {
7076        fn constant(c: f64) -> Self {
7077            <Self as JetScalar<K>>::constant(c)
7078        }
7079        fn add(&self, o: &Self) -> Self {
7080            crate::nested_dual::JetField::add(self, o)
7081        }
7082        fn sub(&self, o: &Self) -> Self {
7083            crate::nested_dual::JetField::sub(self, o)
7084        }
7085        fn mul(&self, o: &Self) -> Self {
7086            crate::nested_dual::JetField::mul(self, o)
7087        }
7088        fn scale(&self, s: f64) -> Self {
7089            crate::nested_dual::JetField::scale(self, s)
7090        }
7091        fn exp(&self) -> Self {
7092            JetScalar::exp(self)
7093        }
7094        fn sqrt(&self) -> Self {
7095            JetScalar::sqrt(self)
7096        }
7097        fn recip(&self) -> Self {
7098            JetScalar::recip(self)
7099        }
7100    }
7101
7102    impl<L: Lane, const K: usize> RowAlg<K> for OneSeedLane<L, K> {
7103        fn constant(c: f64) -> Self {
7104            OneSeedLane::constant(L::splat(c))
7105        }
7106        fn add(&self, o: &Self) -> Self {
7107            OneSeedLane::add(self, o)
7108        }
7109        fn sub(&self, o: &Self) -> Self {
7110            OneSeedLane::sub(self, o)
7111        }
7112        fn mul(&self, o: &Self) -> Self {
7113            OneSeedLane::mul(self, o)
7114        }
7115        fn scale(&self, s: f64) -> Self {
7116            OneSeedLane::scale(self, s)
7117        }
7118        fn exp(&self) -> Self {
7119            OneSeedLane::exp(self)
7120        }
7121        fn sqrt(&self) -> Self {
7122            OneSeedLane::sqrt(self)
7123        }
7124        fn recip(&self) -> Self {
7125            OneSeedLane::recip(self)
7126        }
7127    }
7128
7129    impl<const K: usize> RowAlg<K> for TwoSeed<K> {
7130        fn constant(c: f64) -> Self {
7131            <Self as JetScalar<K>>::constant(c)
7132        }
7133        fn add(&self, o: &Self) -> Self {
7134            crate::nested_dual::JetField::add(self, o)
7135        }
7136        fn sub(&self, o: &Self) -> Self {
7137            crate::nested_dual::JetField::sub(self, o)
7138        }
7139        fn mul(&self, o: &Self) -> Self {
7140            crate::nested_dual::JetField::mul(self, o)
7141        }
7142        fn scale(&self, s: f64) -> Self {
7143            crate::nested_dual::JetField::scale(self, s)
7144        }
7145        fn exp(&self) -> Self {
7146            JetScalar::exp(self)
7147        }
7148        fn sqrt(&self) -> Self {
7149            JetScalar::sqrt(self)
7150        }
7151        fn recip(&self) -> Self {
7152            JetScalar::recip(self)
7153        }
7154    }
7155
7156    impl<L: Lane, const K: usize> RowAlg<K> for TwoSeedLane<L, K> {
7157        fn constant(c: f64) -> Self {
7158            TwoSeedLane::constant(L::splat(c))
7159        }
7160        fn add(&self, o: &Self) -> Self {
7161            TwoSeedLane::add(self, o)
7162        }
7163        fn sub(&self, o: &Self) -> Self {
7164            TwoSeedLane::sub(self, o)
7165        }
7166        fn mul(&self, o: &Self) -> Self {
7167            TwoSeedLane::mul(self, o)
7168        }
7169        fn scale(&self, s: f64) -> Self {
7170            TwoSeedLane::scale(self, s)
7171        }
7172        fn exp(&self) -> Self {
7173            TwoSeedLane::exp(self)
7174        }
7175        fn sqrt(&self) -> Self {
7176            TwoSeedLane::sqrt(self)
7177        }
7178        fn recip(&self) -> Self {
7179            TwoSeedLane::recip(self)
7180        }
7181    }
7182
7183    fn check_oneseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
7184        let mut rows_checked = 0;
7185        for _ in 0..batches {
7186            let rows: [[f64; K]; 4] =
7187                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7188            // Per-row ε-direction.
7189            let u: [[f64; K]; 4] =
7190                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7191
7192            // Production ground truth (scalar OneSeed per row).
7193            let prod: [OneSeed<K>; 4] = std::array::from_fn(|r| {
7194                let p: [OneSeed<K>; K] =
7195                    std::array::from_fn(|a| OneSeed::seed_direction(rows[r][a], a, u[r][a]));
7196                row_expr(&p)
7197            });
7198
7199            // f64-lane re-type per row.
7200            let scal: [OneSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
7201                let p: [OneSeedLane<f64, K>; K] =
7202                    std::array::from_fn(|a| OneSeedLane::seed_direction(rows[r][a], a, u[r][a]));
7203                row_expr(&p)
7204            });
7205
7206            // 4-rows-per-pass batch.
7207            let pbatch: [OneSeedBatch<K>; K] = std::array::from_fn(|a| {
7208                let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7209                let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
7210                OneSeedBatch::seed_direction(val, a, uu)
7211            });
7212            let batch = row_expr(&pbatch);
7213
7214            for r in 0..4 {
7215                let want = prod[r].contracted_third();
7216                let got_scal = scal[r].contracted_third();
7217                let got_batch = batch.lane(r).contracted_third();
7218                // Value channel too (sanity that the base program agrees).
7219                assert_eq!(
7220                    scal[r].base.v.to_bits(),
7221                    prod[r].base.value().to_bits(),
7222                    "OneSeed K={K} scalar value"
7223                );
7224                assert_eq!(
7225                    batch.lane(r).base.value().to_bits(),
7226                    prod[r].base.value().to_bits(),
7227                    "OneSeed K={K} batch lane {r} value"
7228                );
7229                for a in 0..K {
7230                    for b in 0..K {
7231                        assert_eq!(
7232                            got_scal[a][b].to_bits(),
7233                            want[a][b].to_bits(),
7234                            "OneSeed K={K} scalar third[{a}][{b}]"
7235                        );
7236                        assert_eq!(
7237                            got_batch[a][b].to_bits(),
7238                            want[a][b].to_bits(),
7239                            "OneSeed K={K} batch lane {r} third[{a}][{b}]"
7240                        );
7241                    }
7242                }
7243                rows_checked += 1;
7244            }
7245        }
7246        rows_checked
7247    }
7248
7249    fn check_twoseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
7250        let mut rows_checked = 0;
7251        for _ in 0..batches {
7252            let rows: [[f64; K]; 4] =
7253                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7254            let u: [[f64; K]; 4] =
7255                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7256            let v: [[f64; K]; 4] =
7257                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7258
7259            let prod: [TwoSeed<K>; 4] = std::array::from_fn(|r| {
7260                let p: [TwoSeed<K>; K] =
7261                    std::array::from_fn(|a| TwoSeed::seed(rows[r][a], a, u[r][a], v[r][a]));
7262                row_expr(&p)
7263            });
7264
7265            let scal: [TwoSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
7266                let p: [TwoSeedLane<f64, K>; K] =
7267                    std::array::from_fn(|a| TwoSeedLane::seed(rows[r][a], a, u[r][a], v[r][a]));
7268                row_expr(&p)
7269            });
7270
7271            let pbatch: [TwoSeedBatch<K>; K] = std::array::from_fn(|a| {
7272                let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7273                let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
7274                let vv = wide::f64x4::new([v[0][a], v[1][a], v[2][a], v[3][a]]);
7275                TwoSeedBatch::seed(val, a, uu, vv)
7276            });
7277            let batch = row_expr(&pbatch);
7278
7279            for r in 0..4 {
7280                let want = prod[r].contracted_fourth();
7281                let got_scal = scal[r].contracted_fourth();
7282                let got_batch = batch.lane(r).contracted_fourth();
7283                assert_eq!(
7284                    scal[r].base.v.to_bits(),
7285                    prod[r].base.value().to_bits(),
7286                    "TwoSeed K={K} scalar value"
7287                );
7288                assert_eq!(
7289                    batch.lane(r).base.value().to_bits(),
7290                    prod[r].base.value().to_bits(),
7291                    "TwoSeed K={K} batch lane {r} value"
7292                );
7293                for a in 0..K {
7294                    for b in 0..K {
7295                        assert_eq!(
7296                            got_scal[a][b].to_bits(),
7297                            want[a][b].to_bits(),
7298                            "TwoSeed K={K} scalar fourth[{a}][{b}]"
7299                        );
7300                        assert_eq!(
7301                            got_batch[a][b].to_bits(),
7302                            want[a][b].to_bits(),
7303                            "TwoSeed K={K} batch lane {r} fourth[{a}][{b}]"
7304                        );
7305                    }
7306                }
7307                rows_checked += 1;
7308            }
7309        }
7310        rows_checked
7311    }
7312
7313    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: the
7314    /// contracted-third channel of every `OneSeedLane` lane is `to_bits`-identical
7315    /// to the production [`OneSeed`] per row.
7316    #[test]
7317    fn oneseed_lanes_contracted_third_bit_identical() {
7318        let mut state = 0x1234_5678_9ABC_DEF0_u64;
7319        let batches = 2000;
7320        let rows_checked = check_oneseed::<2>(&mut state, batches)
7321            + check_oneseed::<3>(&mut state, batches)
7322            + check_oneseed::<4>(&mut state, batches)
7323            + check_oneseed::<9>(&mut state, batches);
7324        // 4 widths × `batches` batches × 4 rows each: a silently empty inner
7325        // loop would leave this at zero instead of passing as a no-op.
7326        assert_eq!(rows_checked, 4 * batches * 4);
7327    }
7328
7329    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: the
7330    /// contracted-fourth channel of every `TwoSeedLane` lane is `to_bits`-identical
7331    /// to the production [`TwoSeed`] per row.
7332    #[test]
7333    fn twoseed_lanes_contracted_fourth_bit_identical() {
7334        let mut state = 0x0FED_CBA9_8765_4321_u64;
7335        let batches = 2000;
7336        let rows_checked = check_twoseed::<2>(&mut state, batches)
7337            + check_twoseed::<3>(&mut state, batches)
7338            + check_twoseed::<4>(&mut state, batches)
7339            + check_twoseed::<9>(&mut state, batches);
7340        // 4 widths × `batches` batches × 4 rows each: a silently empty inner
7341        // loop would leave this at zero instead of passing as a no-op.
7342        assert_eq!(rows_checked, 4 * batches * 4);
7343    }
7344}
7345
7346#[cfg(test)]
7347mod unit_tests {
7348    use super::{
7349        DynamicJetArena, DynamicOrder2, JetScalar, OneSeed, Order1, Order2, RuntimeJetScalar,
7350        filtered_implicit_solve_scalar,
7351    };
7352    use crate::nested_dual::{Dual2, JetField};
7353
7354    // A deliberately mixed beta/family program. Its family derivatives below
7355    // are analytic scalar programs, so these checks have no finite-difference
7356    // truncation or step-size oracle.
7357    fn family_program<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7358        let xy = x.mul(y);
7359        let exponential = theta.mul(&xy).exp();
7360        let theta_squared_x_squared = theta.mul(theta).mul(&x.mul(x)).scale(0.375);
7361        let theta_y_cubed = theta.mul(&y.mul(y).mul(y)).scale(-0.2);
7362        exponential
7363            .add(&theta_squared_x_squared)
7364            .add(&theta_y_cubed)
7365    }
7366
7367    fn analytic_family_first<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7368        let xy = x.mul(y);
7369        let exponential = theta.mul(&xy).exp();
7370        xy.mul(&exponential)
7371            .add(&theta.mul(&x.mul(x)).scale(0.75))
7372            .add(&y.mul(y).mul(y).scale(-0.2))
7373    }
7374
7375    fn analytic_family_second<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7376        let xy = x.mul(y);
7377        let exponential = theta.mul(&xy).exp();
7378        xy.mul(&xy).mul(&exponential).add(&x.mul(x).scale(0.75))
7379    }
7380
7381    fn assert_channel_close(actual: f64, expected: f64, channel: &str) {
7382        let tolerance = 256.0 * f64::EPSILON * (1.0 + actual.abs().max(expected.abs()));
7383        assert!(
7384            (actual - expected).abs() <= tolerance,
7385            "{channel}: actual={actual:.17e}, expected={expected:.17e}, tolerance={tolerance:.3e}"
7386        );
7387    }
7388
7389    fn assert_order2_channels<const K: usize>(
7390        actual: &Order2<K>,
7391        expected: &Order2<K>,
7392        prefix: &str,
7393    ) {
7394        assert_channel_close(actual.value(), expected.value(), &format!("{prefix}.value"));
7395        for a in 0..K {
7396            assert_channel_close(actual.g()[a], expected.g()[a], &format!("{prefix}.g[{a}]"));
7397            for b in 0..K {
7398                assert_channel_close(
7399                    actual.h()[a][b],
7400                    expected.h()[a][b],
7401                    &format!("{prefix}.h[{a}][{b}]"),
7402                );
7403            }
7404        }
7405    }
7406
7407    /// The outer first/second channels carry the family derivative while each
7408    /// inner `Order2` independently carries beta value/gradient/Hessian.
7409    #[test]
7410    fn dual2_order2_extracts_exact_family_value_gradient_hessian_channels() {
7411        const K: usize = 2;
7412        let x0 = 0.7;
7413        let y0 = -0.45;
7414        let theta0 = 0.6;
7415        let x = <Dual2<Order2<K>> as JetScalar<K>>::variable(x0, 0);
7416        let y = <Dual2<Order2<K>> as JetScalar<K>>::variable(y0, 1);
7417        let theta = Dual2 {
7418            v: Order2::constant(theta0),
7419            g: Order2::constant(1.0),
7420            h: Order2::constant(0.0),
7421        };
7422
7423        let actual = family_program(&x, &y, &theta);
7424        let reference_x = Order2::variable(x0, 0);
7425        let reference_y = Order2::variable(y0, 1);
7426        let reference_theta = Order2::constant(theta0);
7427        let expected_first = analytic_family_first(&reference_x, &reference_y, &reference_theta);
7428        let expected_second = analytic_family_second(&reference_x, &reference_y, &reference_theta);
7429
7430        assert_order2_channels(&actual.g, &expected_first, "family_first");
7431        assert_order2_channels(&actual.h, &expected_second, "family_second");
7432    }
7433
7434    /// `g.eps.h[a][b]` is the directional beta drift of the family-Hessian
7435    /// channel, obtained by exact nested AD rather than differencing Hessians.
7436    #[test]
7437    fn dual2_oneseed_extracts_exact_family_hessian_drift() {
7438        const K: usize = 2;
7439        let x0 = 0.7;
7440        let y0 = -0.45;
7441        let theta0 = 0.6;
7442        let direction = [0.3, -0.8];
7443
7444        let mut x = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(x0, 0);
7445        let mut y = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(y0, 1);
7446        x.v.eps = Order2::constant(direction[0]);
7447        y.v.eps = Order2::constant(direction[1]);
7448        let theta = Dual2 {
7449            v: OneSeed::constant(theta0),
7450            g: OneSeed::constant(1.0),
7451            h: OneSeed::constant(0.0),
7452        };
7453
7454        let actual = family_program(&x, &y, &theta);
7455        let reference_x = OneSeed::seed_direction(x0, 0, direction[0]);
7456        let reference_y = OneSeed::seed_direction(y0, 1, direction[1]);
7457        let reference_theta = OneSeed::constant(theta0);
7458        let expected = analytic_family_first(&reference_x, &reference_y, &reference_theta);
7459
7460        assert_order2_channels(&actual.g.eps, &expected.eps, "family_first_drift");
7461    }
7462
7463    // ── Order2 direct property tests ─────────────────────────────────────────
7464
7465    /// `Order2::constant(c)` carries value `c` and zero everywhere else.
7466    #[test]
7467    fn order2_constant_has_zero_derivatives() {
7468        let s = Order2::<3>::constant(7.5);
7469        assert_eq!(s.value(), 7.5);
7470        for a in 0..3 {
7471            assert_eq!(s.g()[a], 0.0, "grad[{a}] should be zero");
7472            for b in 0..3 {
7473                assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
7474            }
7475        }
7476    }
7477
7478    /// `Order2::variable(x, axis)` has unit gradient in slot `axis` and zero Hessian.
7479    #[test]
7480    fn order2_variable_has_unit_gradient_in_seeded_slot() {
7481        let x = -2.5_f64;
7482        let s = Order2::<4>::variable(x, 2);
7483        assert_eq!(s.value(), x);
7484        for a in 0..4 {
7485            let expected_g = if a == 2 { 1.0 } else { 0.0 };
7486            assert_eq!(s.g()[a], expected_g, "grad[{a}]");
7487            for b in 0..4 {
7488                assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
7489            }
7490        }
7491    }
7492
7493    /// `Order2::add` sums gradient channels; `sub` is the inverse on gradients.
7494    /// Uses integer-valued primaries so the value roundtrip is also exact.
7495    #[test]
7496    fn order2_add_sub_roundtrip() {
7497        let p = Order2::<2>::variable(3.0, 0);
7498        let q = Order2::<2>::variable(2.0, 1);
7499        let pq = crate::nested_dual::JetField::add(&p, &q);
7500        // value = 3 + 2 = 5
7501        assert_eq!(pq.value(), 5.0, "add value");
7502        let back = crate::nested_dual::JetField::sub(&pq, &q);
7503        // (p + q) - q gradient should equal p's gradient exactly
7504        for a in 0..2 {
7505            assert_eq!(back.g()[a], p.g()[a], "grad[{a}] roundtrip");
7506        }
7507    }
7508
7509    /// `Order2::mul` of two variables satisfies the Leibniz product rule:
7510    ///   ∂(p·q)/∂p = q,  ∂(p·q)/∂q = p,  ∂²(p·q)/∂p∂q = 1.
7511    #[test]
7512    fn order2_mul_satisfies_leibniz_rule() {
7513        let pv = 3.0_f64;
7514        let qv = -2.0_f64;
7515        let p = Order2::<2>::variable(pv, 0);
7516        let q = Order2::<2>::variable(qv, 1);
7517        let pq = crate::nested_dual::JetField::mul(&p, &q);
7518        assert_eq!(pq.value(), pv * qv, "value = p·q");
7519        assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
7520        assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
7521        assert_eq!(pq.h()[0][1], 1.0, "∂²(p·q)/∂p∂q = 1");
7522        assert_eq!(pq.h()[1][0], 1.0, "∂²(p·q)/∂q∂p = 1 (symmetric)");
7523        assert_eq!(pq.h()[0][0], 0.0, "∂²(p·q)/∂p² = 0");
7524        assert_eq!(pq.h()[1][1], 0.0, "∂²(p·q)/∂q² = 0");
7525    }
7526
7527    /// `Order2::scale(s)` multiplies every channel by `s`.
7528    #[test]
7529    fn order2_scale_multiplies_all_channels() {
7530        let p = Order2::<2>::variable(4.0, 0);
7531        let s = 2.5_f64;
7532        let ps = crate::nested_dual::JetField::scale(&p, s);
7533        assert_eq!(ps.value(), 4.0 * s);
7534        assert_eq!(ps.g()[0], 1.0 * s);
7535        assert_eq!(ps.g()[1], 0.0);
7536    }
7537
7538    /// `Order2::exp` at a constant has value `e^c`, gradient `e^c * g`, Hessian `e^c * (g⊗g + H)`.
7539    /// At a seeded variable `p₀`, the first derivative is `e^{p₀}` and second is `e^{p₀}`.
7540    #[test]
7541    fn order2_exp_derivative_stack_correct() {
7542        let p0 = 1.0_f64;
7543        let p = Order2::<1>::variable(p0, 0);
7544        let ep = JetScalar::exp(&p);
7545        let e = p0.exp();
7546        assert!((ep.value() - e).abs() < 1e-15, "exp value");
7547        assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p) = exp(p)");
7548        assert!((ep.h()[0][0] - e).abs() < 1e-15, "d²/dp² exp(p) = exp(p)");
7549    }
7550
7551    /// `Order2::ln` at a seeded variable: d/dp ln(p) = 1/p, d²/dp² ln(p) = -1/p².
7552    #[test]
7553    fn order2_ln_derivative_stack_correct() {
7554        let p0 = 2.0_f64;
7555        let p = Order2::<1>::variable(p0, 0);
7556        let lnp = JetScalar::ln(&p);
7557        assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
7558        assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
7559        assert!(
7560            (lnp.h()[0][0] - (-1.0 / (p0 * p0))).abs() < 1e-15,
7561            "d²/dp² ln(p) = -1/p²"
7562        );
7563    }
7564
7565    #[test]
7566    fn dynamic_order2_ln_uses_runtime_scalar_derivative_stack() {
7567        let p0 = 2.0_f64;
7568        let arena = DynamicJetArena::new();
7569        let p = DynamicOrder2::variable(p0, 0, 1, &arena);
7570        let lnp = RuntimeJetScalar::ln(&p);
7571        assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
7572        assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
7573        assert!(
7574            (lnp.h_at(0, 0) - (-1.0 / (p0 * p0))).abs() < 1e-15,
7575            "d²/dp² ln(p) = -1/p²"
7576        );
7577    }
7578
7579    /// `exp` and `ln` are mutual inverses: `ln(exp(p)).value() == p` at the scalar.
7580    #[test]
7581    fn order2_exp_ln_roundtrip_at_value() {
7582        let p0 = 0.8_f64;
7583        let p = Order2::<1>::variable(p0, 0);
7584        let roundtrip = JetScalar::ln(&JetScalar::exp(&p));
7585        assert!((roundtrip.value() - p0).abs() < 1e-14, "ln(exp(p)) ≈ p");
7586    }
7587
7588    // ── Order1 tests ─────────────────────────────────────────────────────────
7589
7590    /// `Order1::constant` carries the correct value with all-zero gradient.
7591    #[test]
7592    fn order1_constant_has_zero_gradient() {
7593        let s = Order1::<3>::constant(-5.0);
7594        assert_eq!(s.value(), -5.0);
7595        for a in 0..3 {
7596            assert_eq!(s.g()[a], 0.0, "g[{a}] should be zero");
7597        }
7598    }
7599
7600    /// `Order1::variable(x, axis)` has unit gradient only in `axis`.
7601    #[test]
7602    fn order1_variable_has_unit_gradient_in_seeded_slot() {
7603        let s = Order1::<3>::variable(2.0, 1);
7604        assert_eq!(s.value(), 2.0);
7605        assert_eq!(s.g()[0], 0.0);
7606        assert_eq!(s.g()[1], 1.0);
7607        assert_eq!(s.g()[2], 0.0);
7608    }
7609
7610    /// `Order1::mul` satisfies the product rule (value and gradient, no Hessian).
7611    #[test]
7612    fn order1_mul_satisfies_product_rule() {
7613        let pv = 3.0_f64;
7614        let qv = -2.0_f64;
7615        let p = Order1::<2>::variable(pv, 0);
7616        let q = Order1::<2>::variable(qv, 1);
7617        let pq = crate::nested_dual::JetField::mul(&p, &q);
7618        assert_eq!(pq.value(), pv * qv);
7619        assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
7620        assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
7621    }
7622
7623    /// `Order1::exp` carries the correct value and gradient `e^{p₀}`.
7624    #[test]
7625    fn order1_exp_has_correct_value_and_gradient() {
7626        let p0 = 0.5_f64;
7627        let p = Order1::<2>::variable(p0, 0);
7628        let ep = JetScalar::exp(&p);
7629        let e = p0.exp();
7630        assert!((ep.value() - e).abs() < 1e-15, "exp value");
7631        assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p)");
7632        assert_eq!(ep.g()[1], 0.0, "irrelevant gradient slot is zero");
7633    }
7634
7635    /// `Order1` and `Order2` agree on value and gradient for the same expression.
7636    #[test]
7637    fn order1_and_order2_agree_on_value_and_gradient() {
7638        let p0 = 1.3_f64;
7639        let q0 = -0.7_f64;
7640        // evaluate (p * q + p).exp() at (p0, q0)
7641        let p1 = Order1::<2>::variable(p0, 0);
7642        let q1 = Order1::<2>::variable(q0, 1);
7643        let expr1 = JetScalar::exp(&crate::nested_dual::JetField::add(
7644            &crate::nested_dual::JetField::mul(&p1, &q1),
7645            &p1,
7646        ));
7647
7648        let p2 = Order2::<2>::variable(p0, 0);
7649        let q2 = Order2::<2>::variable(q0, 1);
7650        let expr2 = JetScalar::exp(&crate::nested_dual::JetField::add(
7651            &crate::nested_dual::JetField::mul(&p2, &q2),
7652            &p2,
7653        ));
7654
7655        assert!(
7656            (expr1.value() - expr2.value()).abs() < 1e-14,
7657            "value mismatch"
7658        );
7659        for a in 0..2 {
7660            assert!(
7661                (expr1.g()[a] - expr2.g()[a]).abs() < 1e-14,
7662                "gradient[{a}] mismatch"
7663            );
7664        }
7665    }
7666
7667    // ── filtered_implicit_solve_scalar ────────────────────────────────────────
7668
7669    /// Lift the trivial linear constraint F(a, θ) = a - θ = 0 through `Order2<1>`.
7670    /// The exact lifted jet is a(θ) = θ, so value=θ₀, gradient=1.
7671    #[test]
7672    fn filtered_implicit_solve_linear_constraint_gives_exact_jet() {
7673        let theta0 = 3.0_f64;
7674        let theta = Order2::<1>::variable(theta0, 0);
7675        // a0 = theta0, F_a = 1, inv_fa = 1; 2 iters suffice for Order2.
7676        let a = filtered_implicit_solve_scalar::<1, Order2<1>>(theta0, 1.0, 2, |a_jet| {
7677            crate::nested_dual::JetField::sub(a_jet, &theta)
7678        });
7679        assert!((a.value() - theta0).abs() < 1e-14, "value = theta0");
7680        // da/dtheta = 1 (identity)
7681        assert!((a.g()[0] - 1.0).abs() < 1e-14, "gradient = 1");
7682        // d²a/dtheta² = 0 (linear)
7683        assert!(a.h()[0][0].abs() < 1e-14, "hessian = 0");
7684    }
7685
7686    /// `filtered_implicit_solve_scalar` on a quadratic constraint F(a,θ)=a²-θ=0
7687    /// with primal root a₀=√θ₀, giving da/dθ = 1/(2√θ₀), d²a/dθ² = -1/(4θ₀^{3/2}).
7688    #[test]
7689    fn filtered_implicit_solve_quadratic_constraint_matches_analytic_derivatives() {
7690        let theta0 = 4.0_f64;
7691        let a0 = theta0.sqrt();
7692        let inv_fa = 1.0 / (2.0 * a0);
7693        let theta = Order2::<1>::variable(theta0, 0);
7694        // F(a,theta) = a*a - theta
7695        let a = filtered_implicit_solve_scalar::<1, Order2<1>>(a0, inv_fa, 2, |a_jet| {
7696            let aa = crate::nested_dual::JetField::mul(a_jet, a_jet);
7697            crate::nested_dual::JetField::sub(&aa, &theta)
7698        });
7699        let tol = 1e-12;
7700        assert!((a.value() - a0).abs() < tol, "value = sqrt(theta0)");
7701        let expected_g = 0.5 / a0;
7702        assert!(
7703            (a.g()[0] - expected_g).abs() < tol,
7704            "da/dtheta = 1/(2*sqrt)"
7705        );
7706        let expected_h = -0.25 / (theta0 * a0);
7707        assert!(
7708            (a.h()[0][0] - expected_h).abs() < tol,
7709            "d2a/dtheta2 = -1/(4*theta^1.5)"
7710        );
7711    }
7712}