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 [`JetScalar::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
53//! [`crate::jet_tower::RowNllProgramGeneric::row_nll_generic`], the seam a family
54//! row loss is written 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/// A truncated-Taylor scalar carrying derivatives in `K` primaries.
66///
67/// All concrete scalars here ([`Order2`], [`OneSeed`], [`TwoSeed`]) and the full
68/// [`crate::jet_tower::Tower4`] implement the SAME algebra; only the carried
69/// channel set differs. A row loss written once against this interface yields a
70/// different channel set per instantiation, all exact for the channel they serve
71/// (doc §A.0).
72pub trait JetScalar<const K: usize>: Copy {
73    /// A constant: value `c`, every derivative channel zero.
74    fn constant(c: f64) -> Self;
75
76    /// The seeded variable `p_axis` at value `x`: unit first derivative in slot
77    /// `axis`, all higher channels zero. (The nilpotent / cross channels of the
78    /// directional scalars are seeded zero — callers set ε/δ directions through
79    /// the scalar-specific [`OneSeed::seed_direction`] / [`TwoSeed::seed`].)
80    fn variable(x: f64, axis: usize) -> Self;
81
82    /// The value channel `ℓ(p)`.
83    fn value(&self) -> f64;
84
85    /// Exact truncated Leibniz sum `self + o`.
86    fn add(&self, o: &Self) -> Self;
87    /// Exact truncated Leibniz difference `self − o`.
88    fn sub(&self, o: &Self) -> Self;
89    /// Exact truncated Leibniz product `self · o`.
90    fn mul(&self, o: &Self) -> Self;
91    /// Negate every channel.
92    fn neg(&self) -> Self;
93    /// Multiply every channel by a plain scalar `s`.
94    fn scale(&self, s: f64) -> Self;
95
96    /// Exact multivariate Faà di Bruno composition `f ∘ self`, given the outer
97    /// derivative stack `d = [f(u), f′(u), f″(u), f‴(u), f⁗(u)]` at
98    /// `u = self.value()`.
99    ///
100    /// This is the SAME `[f64; 5]` stack shape [`crate::jet_tower::Tower4`] and
101    /// the families' `unary_derivatives_*` helpers (built on erfcx / log_ndtr)
102    /// already produce, so those stacks plug in directly. Each scalar consumes
103    /// only the leading entries its order needs (order-2 reads `d[0..=2]`; the
104    /// directional scalars read one / two beyond their base) — the fixed-length
105    /// array makes that windowing total, no length guard required.
106    fn compose_unary(&self, d: [f64; 5]) -> Self;
107
108    /// Compose with a unary special-function whose derivative STACK is built
109    /// from the scalar base value through `stack_fn` — the generic-over-`Lane`
110    /// seam that lets a single-sourced row program instantiate at BOTH the scalar
111    /// `f64` jets and the SIMD `f64x4` batch towers from ONE expression.
112    ///
113    /// On a scalar jet this evaluates `stack_fn(self.value())` ONCE and forwards
114    /// to [`compose_unary`](Self::compose_unary), so it is BIT-IDENTICAL to the
115    /// hand-written `self.compose_unary(stack_fn(self.value()))` (default body
116    /// below). The lever is that the SAME call shape exists on
117    /// [`crate::jet_tower::Tower3Lane`] / [`crate::jet_tower::Tower4Lane`], where
118    /// the four lanes carry FOUR DISTINCT base values, so the batch
119    /// implementation re-runs `stack_fn` per lane — a thing the old
120    /// `compose_unary(stack_from(self.value()))` shape could not express on a
121    /// batch type (it has no single scalar `.value()`). Writing a row program
122    /// against this method instead of the explicit two-step is what makes it
123    /// instantiate, unchanged, at `f64x4` for the 4-rows-per-pass batch path.
124    fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
125        self.compose_unary(stack_fn(self.value()))
126    }
127
128    /// `e^self`. Convenience for tame arguments (see module stability note).
129    fn exp(&self) -> Self {
130        let e = self.value().exp();
131        self.compose_unary([e, e, e, e, e])
132    }
133
134    /// `√self`. Caller guarantees positivity.
135    fn sqrt(&self) -> Self {
136        let u = self.value();
137        let s = u.sqrt();
138        self.compose_unary([
139            s,
140            0.5 / s,
141            -0.25 / (u * s),
142            0.375 / (u * u * s),
143            -0.9375 / (u * u * u * s),
144        ])
145    }
146
147    /// `ln(self)`. Caller guarantees positivity. Same derivative stack
148    /// [`crate::jet_tower::Tower4::ln`] uses, so any program written over both
149    /// matches term-for-term.
150    fn ln(&self) -> Self {
151        let u = self.value();
152        let r = 1.0 / u;
153        self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
154    }
155
156    /// `1/self`.
157    fn recip(&self) -> Self {
158        let r = 1.0 / self.value();
159        let r2 = r * r;
160        self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
161    }
162
163    /// `self^a` for real exponent `a`. Caller guarantees a positive base.
164    /// Mirrors [`crate::jet_tower::Tower4::powf`] (falling-factorial stack).
165    fn powf(&self, a: f64) -> Self {
166        let u = self.value();
167        self.compose_unary([
168            u.powf(a),
169            a * u.powf(a - 1.0),
170            a * (a - 1.0) * u.powf(a - 2.0),
171            a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
172            a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
173        ])
174    }
175
176    /// `ln Γ(self)`. Caller guarantees a positive argument. Uses the SAME
177    /// hand-certified derivative stack [`crate::jet_tower::Tower4::ln_gamma`]
178    /// consumes ([`crate::jet_tower::ln_gamma_derivative_stack`]), so any
179    /// program written over both matches term-for-term.
180    fn ln_gamma(&self) -> Self {
181        self.compose_unary(crate::jet_tower::ln_gamma_derivative_stack(self.value()))
182    }
183
184    /// `ψ(self) = d/dx ln Γ(x)` (digamma). Caller guarantees a positive
185    /// argument. Same hand-certified stack
186    /// [`crate::jet_tower::digamma_derivative_stack`].
187    fn digamma(&self) -> Self {
188        self.compose_unary(crate::jet_tower::digamma_derivative_stack(self.value()))
189    }
190}
191
192/// A Taylor-jet scalar whose primary dimension is selected at runtime.
193///
194/// This is the dimensioned counterpart of [`JetScalar`].  Its algebra is the
195/// same; only the constructors receive the row's actual primary count.  The
196/// fixed-size scalar implementations below bridge to this trait as well, which
197/// lets one row program serve both the const-generic derivative oracles and the
198/// runtime-sized production backends without duplicating the expression.
199pub trait RuntimeJetScalar<'arena>: Clone {
200    /// Storage arena used by runtime-backed scalars. Fixed derivative oracles
201    /// use the unit type because their storage is inline.
202    type Workspace: ?Sized;
203
204    /// A constant in a `dimension`-primary algebra.
205    fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
206    /// A seeded variable in a `dimension`-primary algebra.
207    fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
208    /// Number of primary derivative axes carried by this scalar.
209    fn dimension(&self) -> usize;
210    /// Value channel.
211    fn value(&self) -> f64;
212    /// Exact truncated sum.
213    fn add(&self, o: &Self) -> Self;
214    /// Exact truncated difference.
215    fn sub(&self, o: &Self) -> Self;
216    /// Exact truncated product.
217    fn mul(&self, o: &Self) -> Self;
218    /// Negate every channel.
219    fn neg(&self) -> Self;
220    /// Scale every channel.
221    fn scale(&self, s: f64) -> Self;
222    /// Exact unary composition from the certified derivative stack.
223    fn compose_unary(&self, d: [f64; 5]) -> Self;
224
225    /// `e^self`.
226    fn exp(&self) -> Self {
227        let e = self.value().exp();
228        self.compose_unary([e, e, e, e, e])
229    }
230
231    /// `1/self`.
232    fn recip(&self) -> Self {
233        let r = 1.0 / self.value();
234        let r2 = r * r;
235        self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
236    }
237}
238
239/// Adapter that presents any const-generic [`JetScalar<K>`] through the
240/// runtime-dimension interface.  It is used by derivative oracles so the same
241/// row program can be instantiated at a fixed tower and at a dynamic packed
242/// scalar; production code unwraps the inner fixed tower after evaluation.
243#[derive(Clone, Copy, Debug)]
244pub struct FixedRuntimeJet<S, const K: usize> {
245    inner: S,
246}
247
248impl<S, const K: usize> FixedRuntimeJet<S, K> {
249    /// Recover the wrapped const-generic scalar.
250    #[must_use]
251    pub fn into_inner(self) -> S {
252        self.inner
253    }
254}
255
256impl<'arena, S: JetScalar<K>, const K: usize> RuntimeJetScalar<'arena> for FixedRuntimeJet<S, K> {
257    type Workspace = ();
258
259    fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
260        assert_eq!(dimension, K, "fixed jet dimension mismatch");
261        Self {
262            inner: S::constant(c),
263        }
264    }
265
266    fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
267        assert_eq!(dimension, K, "fixed jet dimension mismatch");
268        Self {
269            inner: S::variable(x, axis),
270        }
271    }
272
273    fn dimension(&self) -> usize {
274        K
275    }
276
277    fn value(&self) -> f64 {
278        self.inner.value()
279    }
280
281    fn add(&self, o: &Self) -> Self {
282        Self {
283            inner: self.inner.add(&o.inner),
284        }
285    }
286
287    fn sub(&self, o: &Self) -> Self {
288        Self {
289            inner: self.inner.sub(&o.inner),
290        }
291    }
292
293    fn mul(&self, o: &Self) -> Self {
294        Self {
295            inner: self.inner.mul(&o.inner),
296        }
297    }
298
299    fn neg(&self) -> Self {
300        Self {
301            inner: self.inner.neg(),
302        }
303    }
304
305    fn scale(&self, s: f64) -> Self {
306        Self {
307            inner: self.inner.scale(s),
308        }
309    }
310
311    fn compose_unary(&self, d: [f64; 5]) -> Self {
312        Self {
313            inner: self.inner.compose_unary(d),
314        }
315    }
316}
317
318/// Reusable storage for runtime-sized packed jets. Scalar primitives write
319/// into this bump arena, so arithmetic performs no heap allocation. Callers
320/// reserve once per worker/chunk and [`reset`](Self::reset) between rows.
321#[derive(Debug)]
322pub struct DynamicJetArena {
323    bump: bumpalo::Bump,
324}
325
326impl DynamicJetArena {
327    /// Create an arena with the allocator's default initial chunk.
328    #[must_use]
329    pub fn new() -> Self {
330        Self {
331            bump: bumpalo::Bump::new(),
332        }
333    }
334
335    /// Create an arena with a row-program-selected initial byte capacity.
336    #[must_use]
337    pub fn with_capacity(bytes: usize) -> Self {
338        Self {
339            bump: bumpalo::Bump::with_capacity(bytes),
340        }
341    }
342
343    /// Reclaim all scalar outputs while retaining allocated chunks.
344    pub fn reset(&mut self) {
345        self.bump.reset();
346    }
347
348    /// Bytes currently reserved from the global allocator. A warm-reset-warm
349    /// benchmark uses this to prove the second row requires no arena growth.
350    #[must_use]
351    pub fn allocated_bytes(&self) -> usize {
352        self.bump.allocated_bytes()
353    }
354
355    #[inline]
356    fn zeros(&self, len: usize) -> &mut [f64] {
357        self.bump.alloc_slice_fill_copy(len, 0.0)
358    }
359
360    /// Allocate and initialize a runtime-sized slice in the arena. Row programs
361    /// use this for their primary-scalar arrays so those arrays share the same
362    /// reusable workspace as derivative channels.
363    pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
364        self.bump.alloc_slice_fill_with(len, fill)
365    }
366}
367
368impl Default for DynamicJetArena {
369    fn default() -> Self {
370        Self::new()
371    }
372}
373
374/// Runtime-sized packed first-order scalar: value plus arena-backed gradient.
375#[derive(Clone, Copy, Debug)]
376pub struct DynamicOrder1<'arena> {
377    arena: &'arena DynamicJetArena,
378    /// Value channel.
379    pub v: f64,
380    /// Gradient channel, length [`Self::dimension`].
381    pub g: &'arena [f64],
382}
383
384impl DynamicOrder1<'_> {
385    /// Gradient channel.
386    #[inline]
387    #[must_use]
388    pub fn g(&self) -> &[f64] {
389        self.g
390    }
391
392    #[inline]
393    fn assert_compatible(&self, o: &Self) {
394        assert_eq!(
395            self.g.len(),
396            o.g.len(),
397            "dynamic first-order jet dimension mismatch"
398        );
399        assert!(
400            std::ptr::eq(self.arena, o.arena),
401            "dynamic jets belong to different arenas"
402        );
403    }
404}
405
406impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder1<'arena> {
407    type Workspace = DynamicJetArena;
408
409    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
410        Self {
411            arena,
412            v: c,
413            g: arena.zeros(dimension),
414        }
415    }
416
417    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
418        assert!(
419            axis < dimension,
420            "dynamic first-order jet axis out of bounds"
421        );
422        let g = arena.zeros(dimension);
423        g[axis] = 1.0;
424        Self { arena, v: x, g }
425    }
426
427    fn dimension(&self) -> usize {
428        self.g.len()
429    }
430    fn value(&self) -> f64 {
431        self.v
432    }
433
434    fn add(&self, o: &Self) -> Self {
435        self.assert_compatible(o);
436        let g = self.arena.zeros(self.dimension());
437        for i in 0..g.len() {
438            g[i] = self.g[i] + o.g[i];
439        }
440        Self {
441            arena: self.arena,
442            v: self.v + o.v,
443            g,
444        }
445    }
446
447    fn sub(&self, o: &Self) -> Self {
448        self.assert_compatible(o);
449        let g = self.arena.zeros(self.dimension());
450        for i in 0..g.len() {
451            g[i] = self.g[i] - o.g[i];
452        }
453        Self {
454            arena: self.arena,
455            v: self.v - o.v,
456            g,
457        }
458    }
459
460    fn mul(&self, o: &Self) -> Self {
461        self.assert_compatible(o);
462        let g = self.arena.zeros(self.dimension());
463        for i in 0..g.len() {
464            g[i] = self.v * o.g[i] + self.g[i] * o.v;
465        }
466        Self {
467            arena: self.arena,
468            v: self.v * o.v,
469            g,
470        }
471    }
472
473    fn neg(&self) -> Self {
474        self.scale(-1.0)
475    }
476
477    fn scale(&self, s: f64) -> Self {
478        let g = self.arena.zeros(self.dimension());
479        for i in 0..g.len() {
480            g[i] = self.g[i] * s;
481        }
482        Self {
483            arena: self.arena,
484            v: self.v * s,
485            g,
486        }
487    }
488
489    fn compose_unary(&self, d: [f64; 5]) -> Self {
490        let g = self.arena.zeros(self.dimension());
491        for i in 0..g.len() {
492            g[i] = d[1] * self.g[i];
493        }
494        Self {
495            arena: self.arena,
496            v: d[0],
497            g,
498        }
499    }
500}
501
502/// Runtime-sized packed second-order scalar: value, gradient, and a row-major
503/// Hessian. Storage is `O(K^2)` in the row's actual primary dimension and comes
504/// from the row's reusable [`DynamicJetArena`].
505#[derive(Clone, Copy, Debug)]
506pub struct DynamicOrder2<'arena> {
507    arena: &'arena DynamicJetArena,
508    /// Value channel.
509    pub v: f64,
510    /// Gradient channel.
511    pub g: &'arena [f64],
512    /// Row-major Hessian channel.
513    pub h: &'arena [f64],
514}
515
516impl DynamicOrder2<'_> {
517    /// Gradient channel.
518    #[inline]
519    #[must_use]
520    pub fn g(&self) -> &[f64] {
521        self.g
522    }
523
524    /// Row-major Hessian channel.
525    #[inline]
526    #[must_use]
527    pub fn h(&self) -> &[f64] {
528        self.h
529    }
530
531    /// Hessian entry `(row, col)`.
532    #[inline]
533    #[must_use]
534    pub fn h_at(&self, row: usize, col: usize) -> f64 {
535        self.h[row * self.dimension() + col]
536    }
537
538    #[inline]
539    fn assert_compatible(&self, o: &Self) {
540        assert_eq!(
541            self.g.len(),
542            o.g.len(),
543            "dynamic second-order jet dimension mismatch"
544        );
545        assert_eq!(
546            self.h.len(),
547            o.h.len(),
548            "dynamic second-order jet Hessian mismatch"
549        );
550        assert!(
551            std::ptr::eq(self.arena, o.arena),
552            "dynamic jets belong to different arenas"
553        );
554    }
555}
556
557impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder2<'arena> {
558    type Workspace = DynamicJetArena;
559
560    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
561        Self {
562            arena,
563            v: c,
564            g: arena.zeros(dimension),
565            h: arena.zeros(dimension * dimension),
566        }
567    }
568
569    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
570        assert!(
571            axis < dimension,
572            "dynamic second-order jet axis out of bounds"
573        );
574        let g = arena.zeros(dimension);
575        g[axis] = 1.0;
576        Self {
577            arena,
578            v: x,
579            g,
580            h: arena.zeros(dimension * dimension),
581        }
582    }
583
584    fn dimension(&self) -> usize {
585        self.g.len()
586    }
587
588    fn value(&self) -> f64 {
589        self.v
590    }
591
592    fn add(&self, o: &Self) -> Self {
593        self.assert_compatible(o);
594        let g = self.arena.zeros(self.dimension());
595        let h = self.arena.zeros(self.h.len());
596        for i in 0..g.len() {
597            g[i] = self.g[i] + o.g[i];
598        }
599        for i in 0..h.len() {
600            h[i] = self.h[i] + o.h[i];
601        }
602        Self {
603            arena: self.arena,
604            v: self.v + o.v,
605            g,
606            h,
607        }
608    }
609
610    fn sub(&self, o: &Self) -> Self {
611        self.assert_compatible(o);
612        let g = self.arena.zeros(self.dimension());
613        let h = self.arena.zeros(self.h.len());
614        for i in 0..g.len() {
615            g[i] = self.g[i] - o.g[i];
616        }
617        for i in 0..h.len() {
618            h[i] = self.h[i] - o.h[i];
619        }
620        Self {
621            arena: self.arena,
622            v: self.v - o.v,
623            g,
624            h,
625        }
626    }
627
628    fn mul(&self, o: &Self) -> Self {
629        self.assert_compatible(o);
630        let n = self.dimension();
631        let g = self.arena.zeros(n);
632        let h = self.arena.zeros(n * n);
633        for i in 0..n {
634            g[i] = self.v * o.g[i] + self.g[i] * o.v;
635        }
636        for i in 0..n {
637            for j in i..n {
638                let ij = i * n + j;
639                let hij =
640                    self.v * o.h[ij] + self.g[i] * o.g[j] + self.g[j] * o.g[i] + self.h[ij] * o.v;
641                h[ij] = hij;
642                h[j * n + i] = hij;
643            }
644        }
645        Self {
646            arena: self.arena,
647            v: self.v * o.v,
648            g,
649            h,
650        }
651    }
652
653    fn neg(&self) -> Self {
654        self.scale(-1.0)
655    }
656
657    fn scale(&self, s: f64) -> Self {
658        let g = self.arena.zeros(self.dimension());
659        let h = self.arena.zeros(self.h.len());
660        for i in 0..g.len() {
661            g[i] = self.g[i] * s;
662        }
663        for i in 0..h.len() {
664            h[i] = self.h[i] * s;
665        }
666        Self {
667            arena: self.arena,
668            v: self.v * s,
669            g,
670            h,
671        }
672    }
673
674    fn compose_unary(&self, d: [f64; 5]) -> Self {
675        let n = self.dimension();
676        let g = self.arena.zeros(n);
677        let h = self.arena.zeros(n * n);
678        for i in 0..n {
679            g[i] = d[1] * self.g[i];
680        }
681        for i in 0..n {
682            for j in 0..n {
683                let ij = i * n + j;
684                h[ij] = d[1] * self.h[ij] + d[2] * self.g[i] * self.g[j];
685            }
686        }
687        Self {
688            arena: self.arena,
689            v: d[0],
690            g,
691            h,
692        }
693    }
694}
695
696/// Runtime-sized one-seed scalar for a Hessian-contracted third derivative.
697#[derive(Clone, Copy, Debug)]
698pub struct DynamicOneSeed<'arena> {
699    /// Base value/gradient/Hessian channels.
700    pub base: DynamicOrder2<'arena>,
701    /// Nilpotent `epsilon` coefficient.
702    pub eps: DynamicOrder2<'arena>,
703}
704
705impl<'arena> DynamicOneSeed<'arena> {
706    /// Seed one primary with the supplied contraction direction component.
707    #[must_use]
708    pub fn seed_direction(
709        x: f64,
710        axis: usize,
711        u_axis: f64,
712        dimension: usize,
713        arena: &'arena DynamicJetArena,
714    ) -> Self {
715        Self {
716            base: DynamicOrder2::variable(x, axis, dimension, arena),
717            eps: DynamicOrder2::constant(u_axis, dimension, arena),
718        }
719    }
720
721    /// Row-major contracted-third matrix.
722    #[must_use]
723    pub fn contracted_third(&self) -> &[f64] {
724        self.eps.h()
725    }
726}
727
728impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeed<'arena> {
729    type Workspace = DynamicJetArena;
730
731    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
732        Self {
733            base: DynamicOrder2::constant(c, dimension, arena),
734            eps: DynamicOrder2::constant(0.0, dimension, arena),
735        }
736    }
737
738    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
739        Self {
740            base: DynamicOrder2::variable(x, axis, dimension, arena),
741            eps: DynamicOrder2::constant(0.0, dimension, arena),
742        }
743    }
744
745    fn dimension(&self) -> usize {
746        self.base.dimension()
747    }
748
749    fn value(&self) -> f64 {
750        self.base.value()
751    }
752
753    fn add(&self, o: &Self) -> Self {
754        Self {
755            base: self.base.add(&o.base),
756            eps: self.eps.add(&o.eps),
757        }
758    }
759
760    fn sub(&self, o: &Self) -> Self {
761        Self {
762            base: self.base.sub(&o.base),
763            eps: self.eps.sub(&o.eps),
764        }
765    }
766
767    fn mul(&self, o: &Self) -> Self {
768        Self {
769            base: self.base.mul(&o.base),
770            eps: self.base.mul(&o.eps).add(&self.eps.mul(&o.base)),
771        }
772    }
773
774    fn neg(&self) -> Self {
775        Self {
776            base: self.base.neg(),
777            eps: self.eps.neg(),
778        }
779    }
780
781    fn scale(&self, s: f64) -> Self {
782        Self {
783            base: self.base.scale(s),
784            eps: self.eps.scale(s),
785        }
786    }
787
788    fn compose_unary(&self, d: [f64; 5]) -> Self {
789        let base = self.base.compose_unary(d);
790        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
791        let eps = fprime.mul(&self.eps);
792        Self { base, eps }
793    }
794}
795
796/// Runtime-sized two-seed scalar for a Hessian-contracted fourth derivative.
797#[derive(Clone, Copy, Debug)]
798pub struct DynamicTwoSeed<'arena> {
799    /// Base value/gradient/Hessian channels.
800    pub base: DynamicOrder2<'arena>,
801    /// Nilpotent `epsilon` coefficient.
802    pub eps: DynamicOrder2<'arena>,
803    /// Nilpotent `delta` coefficient.
804    pub del: DynamicOrder2<'arena>,
805    /// Mixed `epsilon delta` coefficient.
806    pub eps_del: DynamicOrder2<'arena>,
807}
808
809impl<'arena> DynamicTwoSeed<'arena> {
810    /// Seed one primary with both contraction direction components.
811    #[must_use]
812    pub fn seed(
813        x: f64,
814        axis: usize,
815        u_axis: f64,
816        v_axis: f64,
817        dimension: usize,
818        arena: &'arena DynamicJetArena,
819    ) -> Self {
820        Self {
821            base: DynamicOrder2::variable(x, axis, dimension, arena),
822            eps: DynamicOrder2::constant(u_axis, dimension, arena),
823            del: DynamicOrder2::constant(v_axis, dimension, arena),
824            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
825        }
826    }
827
828    /// Row-major contracted-fourth matrix.
829    #[must_use]
830    pub fn contracted_fourth(&self) -> &[f64] {
831        self.eps_del.h()
832    }
833}
834
835impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeed<'arena> {
836    type Workspace = DynamicJetArena;
837
838    fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
839        Self {
840            base: DynamicOrder2::constant(c, dimension, arena),
841            eps: DynamicOrder2::constant(0.0, dimension, arena),
842            del: DynamicOrder2::constant(0.0, dimension, arena),
843            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
844        }
845    }
846
847    fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
848        Self {
849            base: DynamicOrder2::variable(x, axis, dimension, arena),
850            eps: DynamicOrder2::constant(0.0, dimension, arena),
851            del: DynamicOrder2::constant(0.0, dimension, arena),
852            eps_del: DynamicOrder2::constant(0.0, dimension, arena),
853        }
854    }
855
856    fn dimension(&self) -> usize {
857        self.base.dimension()
858    }
859
860    fn value(&self) -> f64 {
861        self.base.value()
862    }
863
864    fn add(&self, o: &Self) -> Self {
865        Self {
866            base: self.base.add(&o.base),
867            eps: self.eps.add(&o.eps),
868            del: self.del.add(&o.del),
869            eps_del: self.eps_del.add(&o.eps_del),
870        }
871    }
872
873    fn sub(&self, o: &Self) -> Self {
874        Self {
875            base: self.base.sub(&o.base),
876            eps: self.eps.sub(&o.eps),
877            del: self.del.sub(&o.del),
878            eps_del: self.eps_del.sub(&o.eps_del),
879        }
880    }
881
882    fn mul(&self, o: &Self) -> Self {
883        let base = self.base.mul(&o.base);
884        let eps = self.base.mul(&o.eps).add(&self.eps.mul(&o.base));
885        let del = self.base.mul(&o.del).add(&self.del.mul(&o.base));
886        let eps_del = self
887            .base
888            .mul(&o.eps_del)
889            .add(&self.eps.mul(&o.del))
890            .add(&self.del.mul(&o.eps))
891            .add(&self.eps_del.mul(&o.base));
892        Self {
893            base,
894            eps,
895            del,
896            eps_del,
897        }
898    }
899
900    fn neg(&self) -> Self {
901        Self {
902            base: self.base.neg(),
903            eps: self.eps.neg(),
904            del: self.del.neg(),
905            eps_del: self.eps_del.neg(),
906        }
907    }
908
909    fn scale(&self, s: f64) -> Self {
910        Self {
911            base: self.base.scale(s),
912            eps: self.eps.scale(s),
913            del: self.del.scale(s),
914            eps_del: self.eps_del.scale(s),
915        }
916    }
917
918    fn compose_unary(&self, d: [f64; 5]) -> Self {
919        let base = self.base.compose_unary(d);
920        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
921        let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]);
922        let eps = fprime.mul(&self.eps);
923        let del = fprime.mul(&self.del);
924        let eps_del = fsecond
925            .mul(&self.eps)
926            .mul(&self.del)
927            .add(&fprime.mul(&self.eps_del));
928        Self {
929            base,
930            eps,
931            del,
932            eps_del,
933        }
934    }
935}
936
937// ── Order2<K> ergonomic operator overloads (doc §A.1) ───────────────────
938//
939// The dispersion-family row NLLs are written with `+`/`-`/`*` operators over
940// the primaries (mirroring how they read as `Tower4` expressions). These
941// delegate channel-for-channel to the inner `Tower2` arithmetic (which has
942// `Add`/`Mul`; `Sub`/`Neg` are expressed as `+ (-1)·rhs` exactly as the
943// `JetScalar::sub` / `JetScalar::neg` impls do), so an `Order2` expression is
944// bit-identical to the same `Tower4` expression's order-≤2 channels.
945
946impl<const K: usize> std::ops::Add for Order2<K> {
947    type Output = Self;
948    #[inline]
949    fn add(self, o: Self) -> Self {
950        Order2(self.0 + o.0)
951    }
952}
953
954impl<const K: usize> std::ops::Add<f64> for Order2<K> {
955    type Output = Self;
956    #[inline]
957    fn add(self, c: f64) -> Self {
958        Order2(self.0 + c)
959    }
960}
961
962impl<const K: usize> std::ops::Sub for Order2<K> {
963    type Output = Self;
964    #[inline]
965    fn sub(self, o: Self) -> Self {
966        Order2(self.0 + o.0.scale(-1.0))
967    }
968}
969
970impl<const K: usize> std::ops::Sub<f64> for Order2<K> {
971    type Output = Self;
972    #[inline]
973    fn sub(self, c: f64) -> Self {
974        Order2(self.0 + (-c))
975    }
976}
977
978impl<const K: usize> std::ops::Mul for Order2<K> {
979    type Output = Self;
980    #[inline]
981    fn mul(self, o: Self) -> Self {
982        Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
983    }
984}
985
986impl<const K: usize> std::ops::Mul<f64> for Order2<K> {
987    type Output = Self;
988    #[inline]
989    fn mul(self, c: f64) -> Self {
990        Order2(self.0.scale(c))
991    }
992}
993
994impl<const K: usize> std::ops::Neg for Order2<K> {
995    type Output = Self;
996    #[inline]
997    fn neg(self) -> Self {
998        Order2(self.0.scale(-1.0))
999    }
1000}
1001
1002/// Filtered Hensel lift of a SCALAR implicit state `a(θ)` defined by the
1003/// constraint `F(a, θ) = 0`, evaluated in ANY [`JetScalar`] algebra `S` (doc
1004/// §11, "A generic implicit-lift operator for every production scalar").
1005///
1006/// This is the perf-respecting alternative to lifting through a dense
1007/// `Tower4<K+1>` (which carries the implicit variable as an extra dense axis):
1008/// the state `a` lives directly in the consumer's own `K`-primary algebra
1009/// `S` — `Order2<K>` for value/gradient/Hessian, `Tower4<K>` for the full
1010/// `t3`/`t4` — never paying for an extra variable.
1011///
1012/// **Method.** Fixed-Jacobian Newton in the nilpotent algebra. By the
1013/// filtered-lift theorem (doc §11.1), if `F_a := ∂F/∂a(a₀, θ₀)` is the primal
1014/// Jacobian at the base point and `inv_fa = 1/F_a`, then the iteration
1015/// `A ← A − inv_fa · F(A, θ)` raises the filtration degree of the residual by
1016/// at least one per step: each step kills exactly one graded layer. Starting
1017/// from `A = const(a₀)` (whose residual lies in `F¹` because `θ − θ₀ ∈ 𝔫`),
1018/// `iters` equal to the algebra's nilpotency order returns the *exact* lifted
1019/// jet (`Order2`: 2, `OneSeed`: 3, `Tower4`/`TwoSeed`: 4). The value channel of
1020/// `A` never moves — `F(A, θ).value() = F(a₀, θ₀) = 0` at the certified root —
1021/// so a caller may precompute every primitive's derivative stack at the fixed
1022/// base index once and let the cheap polynomial composition repeat per step.
1023///
1024/// `f` evaluates the constraint `F(a, θ)` in `S` (capturing the seeded
1025/// parameter jets `θ`); `a0` is the certified scalar root `F(a₀, θ₀) ≈ 0`.
1026pub fn filtered_implicit_solve_scalar<const K: usize, S: JetScalar<K>>(
1027    a0: f64,
1028    inv_fa: f64,
1029    iters: usize,
1030    f: impl Fn(&S) -> S,
1031) -> S {
1032    let mut a = S::constant(a0);
1033    for _ in 0..iters {
1034        let residual = f(&a);
1035        a = a.sub(&residual.scale(inv_fa));
1036    }
1037    a
1038}
1039
1040// ── Order2<K>: value / gradient / Hessian (doc §A.1) ────────────────────
1041
1042/// Truncated SECOND-order scalar: value `v`, gradient `g_a`, Hessian `H_{ab}`.
1043///
1044/// This is a thin newtype over the production [`crate::jet_tower::Tower2`], so
1045/// its `(v, g, H)` channels are obtained by the SAME formulas — and are
1046/// therefore bit-identical to both [`crate::jet_tower::Tower2`] and the order-≤2
1047/// channels of a full [`crate::jet_tower::Tower4`] (doc §A.1, "Bit-identity with
1048/// the full tower"). The wrapper exists only to satisfy the generic
1049/// [`JetScalar`] interface (the `compose_unary` / `add` / `sub` / `neg` /
1050/// `recip` the trait demands, which `Tower2` does not expose by that shape) —
1051/// every channel is delegated to `Tower2` arithmetic unchanged.
1052#[derive(Clone, Copy, Debug)]
1053pub struct Order2<const K: usize>(pub crate::jet_tower::Tower2<K>);
1054
1055impl<const K: usize> Order2<K> {
1056    /// Read the gradient channel `g_a = ∂ℓ/∂p_a`.
1057    #[inline]
1058    pub fn g(&self) -> [f64; K] {
1059        self.0.g
1060    }
1061
1062    /// Read the Hessian channel.
1063    #[inline]
1064    pub fn h(&self) -> [[f64; K]; K] {
1065        self.0.h
1066    }
1067}
1068
1069impl<const K: usize> JetScalar<K> for Order2<K> {
1070    fn constant(c: f64) -> Self {
1071        Order2(crate::jet_tower::Tower2::constant(c))
1072    }
1073    fn variable(x: f64, axis: usize) -> Self {
1074        Order2(crate::jet_tower::Tower2::variable(x, axis))
1075    }
1076    fn value(&self) -> f64 {
1077        self.0.v
1078    }
1079    fn add(&self, o: &Self) -> Self {
1080        Order2(self.0 + o.0)
1081    }
1082    fn sub(&self, o: &Self) -> Self {
1083        // Tower2 has no Sub op; subtract by adding the negation, matching
1084        // Tower4::sub (self + o.scale(-1.0)).
1085        Order2(self.0 + o.0.scale(-1.0))
1086    }
1087    fn mul(&self, o: &Self) -> Self {
1088        Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
1089    }
1090    fn neg(&self) -> Self {
1091        Order2(self.0.scale(-1.0))
1092    }
1093    fn scale(&self, s: f64) -> Self {
1094        Order2(self.0.scale(s))
1095    }
1096    fn compose_unary(&self, d: [f64; 5]) -> Self {
1097        // Order-≤2 reads only [f, f', f''] of the stack.
1098        Order2(self.0.compose_unary([d[0], d[1], d[2]]))
1099    }
1100}
1101
1102// ── Lane-batched Order-2 scalar: 4 rows per pass in SIMD lanes (perf) ────
1103//
1104// The hot per-row jet kernels evaluate ONE row's `(v, g, H)` tower at a time in
1105// scalar `f64`. A hand-written scalar derivative does the same. The throughput
1106// lever a jet has that scalar hand-code cannot is **row batching in SIMD
1107// lanes**: the order-≤2 Leibniz product `Order2::mul` is `O(K²)` independent
1108// per-channel float ops, and EVERY row runs the identical op graph on different
1109// data — the textbook SPMD shape. Packing `LANES = 4` rows into a `wide::f64x4`
1110// and running the algebra once per 4 rows replaces 4 scalar passes with one
1111// vector pass: the `K²` Hessian channel updates become `K²` NEON `.2d` / SSE2
1112// `pd` instructions covering 4 rows each, ~4× fewer FP instructions per row.
1113//
1114// The carried scalar field is abstracted by [`Lane`] so the SAME algebra body
1115// instantiates at `f64` (1 row, used as the bit-identity oracle) or
1116// [`wide::f64x4`] (4 rows). Bit-identity is structural, not approximate:
1117//
1118//   * Every arithmetic op is a plain lane-wise `+` / `-` / `*` (NEVER a fused
1119//     `mul_add`), and IEEE-754 double `+`/`-`/`*`/`/` are correctly rounded and
1120//     deterministic, so lane `i` of an `f64x4` op equals the scalar `f64` op on
1121//     that lane's inputs bit-for-bit.
1122//   * The transcendental derivative STACKS (`exp`/`ln`/`sqrt`/…) are produced
1123//     **per lane by the identical scalar code** ([`Lane::unary3`] unpacks, runs
1124//     the same `[f64; 3]` stack closure the scalar path runs, repacks), so the
1125//     only thing vectorised is the cheap rational tensor composition — the
1126//     library transcendental itself is the exact same `f64::exp` call per lane.
1127//   * The op order mirrors [`crate::jet_tower::Tower2`] term-for-term, so
1128//     [`Order2Lane<f64, K>`] is `to_bits`-identical to the production
1129//     [`Order2<K>`] (= `Tower2<K>`), and [`Order2Lane<f64x4, K>`] lane `i` is
1130//     `to_bits`-identical to that — proven by the `batch_tests` oracle below
1131//     (≥2000 random 4-row batches across `K ∈ {2,3,4,9}`).
1132
1133/// The scalar field a [`Order2Lane`] carries: either a single `f64` (one row,
1134/// the oracle) or a [`wide::f64x4`] (four rows evaluated in SIMD lanes). All ops
1135/// are plain lane-wise IEEE arithmetic, so a vector op equals the scalar op on
1136/// each lane bit-for-bit.
1137pub trait Lane: Copy {
1138    /// Broadcast a scalar to every lane.
1139    fn splat(x: f64) -> Self;
1140    /// Lane-wise `self + o`.
1141    fn add(self, o: Self) -> Self;
1142    /// Lane-wise `self - o`.
1143    fn sub(self, o: Self) -> Self;
1144    /// Lane-wise `self * o`.
1145    fn mul(self, o: Self) -> Self;
1146    /// The `f64` in lane `i` (`i < LANES`; `f64` ignores `i`).
1147    fn lane(self, i: usize) -> f64;
1148    /// Build the order-≤2 derivative stack `[f(u), f′(u), f″(u)]` **per lane**
1149    /// from the lane value `u`, via the SAME scalar `stack` closure the
1150    /// per-row path runs (so the transcendental/rational stack is bit-identical
1151    /// to the scalar evaluation — only the subsequent tensor composition is
1152    /// vectorised).
1153    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3];
1154    /// Build the order-≤4 derivative stack `[f, f′, f″, f‴, f⁗]` **per lane**
1155    /// from the lane value `u`, via the SAME scalar `stack` closure the per-row
1156    /// path runs. The one-/two-seed scalars ([`OneSeedLane`] / [`TwoSeedLane`])
1157    /// need outer derivatives one / two orders beyond their order-2 base, so
1158    /// they build their composition stack through this five-entry variant. As
1159    /// with [`unary3`](Lane::unary3), only the transcendental/rational stack is
1160    /// evaluated per lane (bit-identically to the scalar path); the subsequent
1161    /// tensor composition is vectorised.
1162    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5];
1163    /// The general-`N` sibling of [`unary3`](Lane::unary3) / [`unary5`](Lane::unary5):
1164    /// build an `N`-wide derivative stack **per lane** from the lane value, via
1165    /// the SAME scalar `stack` closure the per-row path runs, then pack the `N`
1166    /// columns lane-wise. This is the lane primitive the compose-with-stack seam
1167    /// ([`crate::jet_tower::Tower4Lane::compose_unary_with`] and its `Tower3`
1168    /// sibling) routes through: it evaluates `stack` once per lane at that lane's
1169    /// OWN base value (each of the four rows in an `f64x4` carries a distinct
1170    /// base), so lane `i` of the packed result equals the scalar `stack(value_i)`
1171    /// bit-for-bit (only the cheap pack is vectorised; the closure body is the
1172    /// identical scalar code). With `N = 3` / `N = 5` it is `to_bits`-identical to
1173    /// [`unary3`](Lane::unary3) / [`unary5`](Lane::unary5).
1174    fn unary_with<const N: usize>(self, stack: impl Fn(f64) -> [f64; N]) -> [Self; N];
1175}
1176
1177impl Lane for f64 {
1178    #[inline]
1179    fn splat(x: f64) -> Self {
1180        x
1181    }
1182    #[inline]
1183    fn add(self, o: Self) -> Self {
1184        self + o
1185    }
1186    #[inline]
1187    fn sub(self, o: Self) -> Self {
1188        self - o
1189    }
1190    #[inline]
1191    fn mul(self, o: Self) -> Self {
1192        self * o
1193    }
1194    #[inline]
1195    fn lane(self, _: usize) -> f64 {
1196        self
1197    }
1198    #[inline]
1199    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
1200        stack(self)
1201    }
1202    #[inline]
1203    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
1204        stack(self)
1205    }
1206    #[inline]
1207    fn unary_with<const N: usize>(self, stack: impl Fn(f64) -> [f64; N]) -> [Self; N] {
1208        // One row: the packed result IS the scalar stack ([Self; N] = [f64; N]).
1209        stack(self)
1210    }
1211}
1212
1213impl Lane for wide::f64x4 {
1214    #[inline]
1215    fn splat(x: f64) -> Self {
1216        wide::f64x4::splat(x)
1217    }
1218    #[inline]
1219    fn add(self, o: Self) -> Self {
1220        self + o
1221    }
1222    #[inline]
1223    fn sub(self, o: Self) -> Self {
1224        self - o
1225    }
1226    #[inline]
1227    fn mul(self, o: Self) -> Self {
1228        self * o
1229    }
1230    #[inline]
1231    fn lane(self, i: usize) -> f64 {
1232        self.to_array()[i]
1233    }
1234    #[inline]
1235    fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
1236        let a = self.to_array();
1237        let mut d0 = [0.0_f64; 4];
1238        let mut d1 = [0.0_f64; 4];
1239        let mut d2 = [0.0_f64; 4];
1240        for i in 0..4 {
1241            let s = stack(a[i]);
1242            d0[i] = s[0];
1243            d1[i] = s[1];
1244            d2[i] = s[2];
1245        }
1246        [
1247            wide::f64x4::new(d0),
1248            wide::f64x4::new(d1),
1249            wide::f64x4::new(d2),
1250        ]
1251    }
1252    #[inline]
1253    fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
1254        let a = self.to_array();
1255        let mut d = [[0.0_f64; 4]; 5];
1256        for i in 0..4 {
1257            let s = stack(a[i]);
1258            for (k, dk) in d.iter_mut().enumerate() {
1259                dk[i] = s[k];
1260            }
1261        }
1262        [
1263            wide::f64x4::new(d[0]),
1264            wide::f64x4::new(d[1]),
1265            wide::f64x4::new(d[2]),
1266            wide::f64x4::new(d[3]),
1267            wide::f64x4::new(d[4]),
1268        ]
1269    }
1270    #[inline]
1271    fn unary_with<const N: usize>(self, stack: impl Fn(f64) -> [f64; N]) -> [Self; N] {
1272        // Evaluate the scalar stack PER LANE at that lane's own base value, then
1273        // pack the N derivative columns lane-wise (the same shape `unary5` uses,
1274        // generalised to N). Lane `i` of column `k` is `stack(base_i)[k]`.
1275        let a = self.to_array();
1276        let mut cols = [[0.0_f64; 4]; N];
1277        for (i, &base) in a.iter().enumerate() {
1278            let s = stack(base);
1279            for (k, sk) in s.iter().enumerate() {
1280                cols[k][i] = *sk;
1281            }
1282        }
1283        std::array::from_fn(|k| wide::f64x4::new(cols[k]))
1284    }
1285}
1286
1287/// A lane-batched order-≤2 Taylor scalar: value / gradient / Hessian carried in
1288/// a SIMD field [`L: Lane`](Lane). With `L = f64x4` one instance carries FOUR
1289/// rows at once, so the row loop processes 4 rows per vector pass instead of one
1290/// per scalar pass.
1291///
1292/// The channel layout and every float op mirror [`crate::jet_tower::Tower2`]
1293/// term-for-term, so `Order2Lane<f64, K>` is `to_bits`-identical to the
1294/// production [`Order2<K>`] and `Order2Lane<f64x4, K>` lane `i` is
1295/// `to_bits`-identical to that (see the module note and `batch_tests`).
1296#[derive(Clone, Copy, Debug)]
1297pub struct Order2Lane<L: Lane, const K: usize> {
1298    /// Value channel `ℓ` (one entry per lane/row).
1299    pub v: L,
1300    /// Gradient channel `∂ℓ/∂p_a`.
1301    pub g: [L; K],
1302    /// Hessian channel `∂²ℓ/∂p_a∂p_b` (symmetric).
1303    pub h: [[L; K]; K],
1304}
1305
1306/// The 4-rows-per-pass batched order-≤2 scalar (`wide::f64x4` lanes).
1307pub type Order2Batch<const K: usize> = Order2Lane<wide::f64x4, K>;
1308
1309impl<L: Lane, const K: usize> Order2Lane<L, K> {
1310    /// A constant: value `c` in every channel-zero slot.
1311    #[inline]
1312    pub fn constant(c: L) -> Self {
1313        Order2Lane {
1314            v: c,
1315            g: [L::splat(0.0); K],
1316            h: [[L::splat(0.0); K]; K],
1317        }
1318    }
1319
1320    /// The seeded variable `p_axis` at (per-lane) value `value`: unit first
1321    /// derivative in slot `axis`. With `L = f64x4`, `value` packs the four
1322    /// rows' values of primary `axis`.
1323    #[inline]
1324    pub fn variable(value: L, axis: usize) -> Self {
1325        let mut out = Self::constant(value);
1326        out.g[axis] = L::splat(1.0);
1327        out
1328    }
1329
1330    /// Lane-wise `self + o` (mirrors `Tower2` Add: per-channel add).
1331    #[inline]
1332    pub fn add(&self, o: &Self) -> Self {
1333        let mut out = *self;
1334        out.v = self.v.add(o.v);
1335        for i in 0..K {
1336            out.g[i] = self.g[i].add(o.g[i]);
1337            for j in 0..K {
1338                out.h[i][j] = self.h[i][j].add(o.h[i][j]);
1339            }
1340        }
1341        out
1342    }
1343
1344    /// Multiply every channel by the plain scalar `s` (mirrors `Tower2::scale`).
1345    #[inline]
1346    pub fn scale(&self, s: f64) -> Self {
1347        let sl = L::splat(s);
1348        let mut out = *self;
1349        out.v = self.v.mul(sl);
1350        for i in 0..K {
1351            out.g[i] = self.g[i].mul(sl);
1352            for j in 0..K {
1353                out.h[i][j] = self.h[i][j].mul(sl);
1354            }
1355        }
1356        out
1357    }
1358
1359    /// Lane-wise `self - o`, expressed as `self + o·(-1)` exactly as
1360    /// [`Order2::sub`] / `Tower4::sub` do, so signed-zero handling matches.
1361    #[inline]
1362    pub fn sub(&self, o: &Self) -> Self {
1363        self.add(&o.scale(-1.0))
1364    }
1365
1366    /// Negate every channel (= `scale(-1.0)`, matching [`Order2::neg`]).
1367    #[inline]
1368    pub fn neg(&self) -> Self {
1369        self.scale(-1.0)
1370    }
1371
1372    /// Exact order-≤2 Leibniz product, term-for-term identical to
1373    /// [`crate::jet_tower::Tower2::mul`] (same factor order, no `mul_add`).
1374    ///
1375    /// The Hessian channel is symmetric under `i ↔ j` (see
1376    /// [`crate::jet_tower::Tower2::mul`] for why the invariant always holds), so
1377    /// we compute the upper triangle `j ≥ i` and mirror it — `K(K+1)/2` lane
1378    /// entry-chains instead of `K²`. Because each lane entry is already a full
1379    /// SIMD op (no cross-`j` lane packing to lose), halving the entry count is a
1380    /// direct throughput win (~18 % on `Order2Lane<f64x4, 9>`, the survival batch
1381    /// kernel, and ~2× on the `f64` oracle). The upper triangle uses the EXACT
1382    /// term order of `Tower2::mul`, so `Order2Lane<f64>` stays `to_bits`-identical
1383    /// to `Order2` (= `Tower2`) and `Order2Lane<f64x4>` lane `i` stays
1384    /// `to_bits`-identical to that; the mirror makes the batch Hessian exactly
1385    /// symmetric, matching the scalar `Tower2::mul` (which mirrors identically).
1386    #[inline]
1387    pub fn mul(&self, o: &Self) -> Self {
1388        let a = self;
1389        let b = o;
1390        let mut out = Self::constant(a.v.mul(b.v));
1391        for i in 0..K {
1392            // a.v*b.g[i] + a.g[i]*b.v
1393            out.g[i] = a.v.mul(b.g[i]).add(a.g[i].mul(b.v));
1394        }
1395        for i in 0..K {
1396            for j in i..K {
1397                // a.v*b.h + a.g[i]*b.g[j] + a.g[j]*b.g[i] + a.h*b.v
1398                let hij =
1399                    a.v.mul(b.h[i][j])
1400                        .add(a.g[i].mul(b.g[j]))
1401                        .add(a.g[j].mul(b.g[i]))
1402                        .add(a.h[i][j].mul(b.v));
1403                out.h[i][j] = hij;
1404                out.h[j][i] = hij;
1405            }
1406        }
1407        out
1408    }
1409
1410    /// Exact order-≤2 Faà di Bruno composition `f ∘ self`, given the per-lane
1411    /// derivative stack `d = [f(u), f′(u), f″(u)]`. Mirrors
1412    /// [`crate::jet_tower::Tower2::compose_unary`] term-for-term (`acc` starts at
1413    /// `0` then accumulates, so signed-zero collapses identically).
1414    #[inline]
1415    pub fn compose_unary(&self, d: [L; 3]) -> Self {
1416        let mut out = Self::constant(d[0]);
1417        for i in 0..K {
1418            let mut acc = L::splat(0.0);
1419            acc = acc.add(d[1].mul(self.g[i]));
1420            out.g[i] = acc;
1421        }
1422        for i in 0..K {
1423            for j in 0..K {
1424                let mut acc = L::splat(0.0);
1425                acc = acc.add(d[1].mul(self.h[i][j]));
1426                acc = acc.add(d[2].mul(self.g[i]).mul(self.g[j]));
1427                out.h[i][j] = acc;
1428            }
1429        }
1430        out
1431    }
1432
1433    /// `e^self`, per-lane stack `[e, e, e]` (matches the [`JetScalar::exp`]
1434    /// default forwarded through `Order2`).
1435    #[inline]
1436    pub fn exp(&self) -> Self {
1437        let d = self.v.unary3(|u| {
1438            let e = u.exp();
1439            [e, e, e]
1440        });
1441        self.compose_unary(d)
1442    }
1443
1444    /// `ln(self)`; caller guarantees positivity. Per-lane stack
1445    /// `[ln u, 1/u, -1/u²]` (matches [`JetScalar::ln`] truncated to order 2).
1446    #[inline]
1447    pub fn ln(&self) -> Self {
1448        let d = self.v.unary3(|u| {
1449            let r = 1.0 / u;
1450            [u.ln(), r, -r * r]
1451        });
1452        self.compose_unary(d)
1453    }
1454
1455    /// `√self`; caller guarantees positivity. Per-lane stack
1456    /// `[s, 0.5/s, -0.25/(u·s)]` (matches [`JetScalar::sqrt`]).
1457    #[inline]
1458    pub fn sqrt(&self) -> Self {
1459        let d = self.v.unary3(|u| {
1460            let s = u.sqrt();
1461            [s, 0.5 / s, -0.25 / (u * s)]
1462        });
1463        self.compose_unary(d)
1464    }
1465
1466    /// `1/self`. Per-lane stack `[r, -r², 2r³]` (matches [`JetScalar::recip`]).
1467    #[inline]
1468    pub fn recip(&self) -> Self {
1469        let d = self.v.unary3(|u| {
1470            let r = 1.0 / u;
1471            let r2 = r * r;
1472            [r, -r2, 2.0 * r2 * r]
1473        });
1474        self.compose_unary(d)
1475    }
1476
1477    /// `self^a` for real `a`; caller guarantees a positive base. Per-lane
1478    /// falling-factorial stack (matches [`JetScalar::powf`]).
1479    #[inline]
1480    pub fn powf(&self, a: f64) -> Self {
1481        let d = self.v.unary3(|u| {
1482            [
1483                u.powf(a),
1484                a * u.powf(a - 1.0),
1485                a * (a - 1.0) * u.powf(a - 2.0),
1486            ]
1487        });
1488        self.compose_unary(d)
1489    }
1490}
1491
1492impl<const K: usize> Order2Batch<K> {
1493    /// Extract lane `i`'s `(v, g, H)` as a production [`Order2<K>`] scalar.
1494    /// Lane `i` is `to_bits`-identical to evaluating the same program at
1495    /// [`Order2<K>`] on row `i` (see `batch_tests`).
1496    #[inline]
1497    #[must_use]
1498    pub fn lane(&self, i: usize) -> Order2<K> {
1499        let mut t = crate::jet_tower::Tower2::<K>::constant(self.v.lane(i));
1500        for a in 0..K {
1501            t.g[a] = self.g[a].lane(i);
1502            for b in 0..K {
1503                t.h[a][b] = self.h[a][b].lane(i);
1504            }
1505        }
1506        Order2(t)
1507    }
1508}
1509
1510// ── Order1<K>: value / gradient only (doc §A.1, first-order prune) ──────
1511
1512/// Truncated FIRST-order scalar: value `v` and gradient `g_a` only — NO Hessian.
1513///
1514/// This is [`Order2`] with the K×K Hessian channel deleted. Its value and
1515/// gradient are computed by the SAME order-≤1 truncation of the Leibniz / Faà
1516/// di Bruno rules that [`Order2`] uses for those two channels, with the float
1517/// operations applied in the identical order — so its `(v, g)` is BIT-IDENTICAL
1518/// to both [`Order2`]'s and a full [`crate::jet_tower::Tower4`]'s order-≤1
1519/// channels. Use it at a consumer that reads ONLY value + gradient (the SAE
1520/// β-border channel: the reconstruction is linear in β, so the Hessian-in-β
1521/// vanishes and the dense K×K Hessian product `Tower2::mul` would build is pure
1522/// discarded work). Order-≤1 value/gradient never read any input's Hessian, so
1523/// dropping that channel changes neither result nor float-op order — it only
1524/// removes the `K²` arithmetic that produced an unread tensor.
1525#[derive(Clone, Copy, Debug)]
1526pub struct Order1<const K: usize> {
1527    /// Value ℓ.
1528    pub v: f64,
1529    /// Gradient ∂ℓ/∂p_a.
1530    pub g: [f64; K],
1531}
1532
1533impl<const K: usize> Order1<K> {
1534    /// Read the gradient channel `g_a = ∂ℓ/∂p_a`.
1535    #[inline]
1536    pub fn g(&self) -> [f64; K] {
1537        self.g
1538    }
1539}
1540
1541impl<const K: usize> JetScalar<K> for Order1<K> {
1542    fn constant(c: f64) -> Self {
1543        // Order2::constant -> Tower2::constant: value c, all derivatives zero.
1544        Order1 { v: c, g: [0.0; K] }
1545    }
1546    fn variable(x: f64, axis: usize) -> Self {
1547        // Order2::variable -> Tower2::variable: unit first derivative in `axis`.
1548        let mut g = [0.0; K];
1549        g[axis] = 1.0;
1550        Order1 { v: x, g }
1551    }
1552    fn value(&self) -> f64 {
1553        self.v
1554    }
1555    fn add(&self, o: &Self) -> Self {
1556        // Tower2 Add: out.v += o.v; out.g[i] += o.g[i] (same float order).
1557        let mut g = self.g;
1558        for i in 0..K {
1559            g[i] += o.g[i];
1560        }
1561        Order1 { v: self.v + o.v, g }
1562    }
1563    fn sub(&self, o: &Self) -> Self {
1564        // Mirror Order2::sub == self + o.scale(-1.0) exactly: scale then add.
1565        self.add(&o.scale(-1.0))
1566    }
1567    fn mul(&self, o: &Self) -> Self {
1568        // Tower2::mul value/grad terms, identical float order:
1569        //   v = a.v*b.v;  g[i] = a.v*b.g[i] + a.g[i]*b.v.
1570        // (The Hessian loop `a.v*b.h + a.g*b.g + ... + a.h*b.v` is the discarded
1571        //  work this type exists to skip; it never feeds v or g.)
1572        let a = self;
1573        let b = o;
1574        let mut g = [0.0; K];
1575        for i in 0..K {
1576            g[i] = a.v * b.g[i] + a.g[i] * b.v;
1577        }
1578        Order1 { v: a.v * b.v, g }
1579    }
1580    fn neg(&self) -> Self {
1581        // Order2::neg == self.0.scale(-1.0).
1582        self.scale(-1.0)
1583    }
1584    fn scale(&self, s: f64) -> Self {
1585        // Tower2::scale: out.v *= s; out.g[i] *= s (same float order).
1586        let mut g = self.g;
1587        for i in 0..K {
1588            g[i] *= s;
1589        }
1590        Order1 { v: self.v * s, g }
1591    }
1592    fn compose_unary(&self, d: [f64; 5]) -> Self {
1593        // Faà di Bruno truncated to order ≤ 1 (matches `faa_di_bruno` /
1594        // `Tower2::compose_unary` for the value and gradient channels):
1595        //   value channel (m=0): d[0].
1596        //   grad channel (positions=[i], single partition {{0}}): d[1]·g[i].
1597        // Order-≤1 reads only d[0], d[1]; trailing stack entries are unused.
1598        let mut g = [0.0; K];
1599        for i in 0..K {
1600            g[i] = d[1] * self.g[i];
1601        }
1602        Order1 { v: d[0], g }
1603    }
1604}
1605
1606// ── OneSeed<K>: one-seed directional, contracted third (doc §A.2) ───────
1607
1608/// One-seed directional scalar: an [`Order2`] base plus ONE nilpotent ε
1609/// (`ε² = 0`) whose coefficient is itself an [`Order2`].
1610///
1611/// A scalar is `s = base + ε·eps`. Arithmetic is the `ε² = 0` truncation of the
1612/// product (doc §A.2): the base parts multiply as ordinary [`Order2`] products,
1613/// and the ε-coefficient picks up `a.base·b.eps + a.eps·b.base`. Composition
1614/// pushes ε through one extra outer derivative.
1615///
1616/// Seed each primary with [`seed_direction`](Self::seed_direction): the base is
1617/// the usual seeded variable (carrying `e_a` for the Hessian channel) and the
1618/// ε-coefficient is the FIXED contraction direction `u_a` (a constant). Then the
1619/// ε-component of the evaluated Hessian channel is the contracted third
1620/// `[eps.h][a][b] = Σ_c ℓ_{abc} u_c` — exactly `row_third_contracted(dir = u)`,
1621/// without materialising `t3`.
1622#[derive(Clone, Copy, Debug)]
1623pub struct OneSeed<const K: usize> {
1624    /// The `ε⁰` part: value / gradient / Hessian of `ℓ`.
1625    pub base: Order2<K>,
1626    /// The `ε¹` part: value / gradient / Hessian of the ε-coefficient. After a
1627    /// `seed_direction(u)` evaluation, `eps.h[a][b] = Σ_c ℓ_{abc} u_c`.
1628    pub eps: Order2<K>,
1629}
1630
1631impl<const K: usize> OneSeed<K> {
1632    /// Seed primary `axis` at value `x` with ε-direction component `u_axis`:
1633    /// `p_axis = p_axis⁰ + x-seed + ε·u_axis`, i.e. base = `variable(x, axis)`
1634    /// and eps = `constant(u_axis)` (doc §A.2 "Seeding").
1635    pub fn seed_direction(x: f64, axis: usize, u_axis: f64) -> Self {
1636        OneSeed {
1637            base: Order2::variable(x, axis),
1638            eps: Order2::constant(u_axis),
1639        }
1640    }
1641
1642    /// The contracted-third channel after a `seed_direction(u)` evaluation:
1643    /// `out[a][b] = Σ_c ℓ_{abc} u_c`, i.e. the ε-coefficient's Hessian (doc §A.2).
1644    pub fn contracted_third(&self) -> [[f64; K]; K] {
1645        self.eps.h()
1646    }
1647}
1648
1649impl<const K: usize> JetScalar<K> for OneSeed<K> {
1650    fn constant(c: f64) -> Self {
1651        OneSeed {
1652            base: Order2::constant(c),
1653            eps: Order2::constant(0.0),
1654        }
1655    }
1656    fn variable(x: f64, axis: usize) -> Self {
1657        // No ε-direction unless seeded via `seed_direction`.
1658        OneSeed {
1659            base: Order2::variable(x, axis),
1660            eps: Order2::constant(0.0),
1661        }
1662    }
1663    fn value(&self) -> f64 {
1664        self.base.value()
1665    }
1666    fn add(&self, o: &Self) -> Self {
1667        OneSeed {
1668            base: self.base.add(&o.base),
1669            eps: self.eps.add(&o.eps),
1670        }
1671    }
1672    fn sub(&self, o: &Self) -> Self {
1673        OneSeed {
1674            base: self.base.sub(&o.base),
1675            eps: self.eps.sub(&o.eps),
1676        }
1677    }
1678    fn mul(&self, o: &Self) -> Self {
1679        // (a.base + ε a.eps)(b.base + ε b.eps), dropping ε².
1680        OneSeed {
1681            base: self.base.mul(&o.base),
1682            eps: self.base.mul(&o.eps).add(&self.eps.mul(&o.base)),
1683        }
1684    }
1685    fn neg(&self) -> Self {
1686        OneSeed {
1687            base: self.base.neg(),
1688            eps: self.eps.neg(),
1689        }
1690    }
1691    fn scale(&self, s: f64) -> Self {
1692        OneSeed {
1693            base: self.base.scale(s),
1694            eps: self.eps.scale(s),
1695        }
1696    }
1697    fn compose_unary(&self, d: [f64; 5]) -> Self {
1698        // f(base + ε eps) = f(base) + ε · f'(base)·eps  (ε² = 0). Each factor is
1699        // an Order2 composition: the base composes with the f-stack, and the
1700        // ε-coefficient is the Order2 of the SHIFTED stack (the chain rule
1701        // `f'(base)` as an Order2) times eps. Order2 reads only the leading
1702        // three entries of whatever stack it is handed, so the trailing slots
1703        // are unused padding (the fixed-length array makes the windowing total).
1704        let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
1705        // f'(base) as an Order2 (consumes [f', f'', f''']).
1706        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
1707        let eps = fprime.mul(&self.eps);
1708        OneSeed { base, eps }
1709    }
1710}
1711
1712// ── OneSeedLane<L, K>: lane-batched one-seed directional (doc §A.2) ──────
1713
1714/// Lane-batched [`OneSeed`]: the same one-seed directional scalar with its two
1715/// [`Order2`] parts re-typed to [`Order2Lane<L, K>`], so one `L = f64x4`
1716/// instance carries FOUR rows' contracted-third evaluations per vector pass.
1717///
1718/// Every operation (`add`/`sub`/`mul`/`neg`/`scale`/`compose_unary` and the
1719/// transcendentals) is a term-for-term structural re-type of the scalar
1720/// [`OneSeed`] ops onto the lane-implemented [`Order2Lane`] algebra. With
1721/// `L = f64`, `OneSeedLane<f64, K>` is `to_bits`-identical to [`OneSeed<K>`];
1722/// with `L = f64x4`, lane `i` is `to_bits`-identical to that (see `batch_tests`).
1723#[derive(Clone, Copy, Debug)]
1724pub struct OneSeedLane<L: Lane, const K: usize> {
1725    /// The `ε⁰` part (lane-batched value / gradient / Hessian of `ℓ`).
1726    pub base: Order2Lane<L, K>,
1727    /// The `ε¹` part. After a `seed_direction(u)` evaluation,
1728    /// `eps.h[a][b]` lane `i` is row `i`'s `Σ_c ℓ_{abc} u_c`.
1729    pub eps: Order2Lane<L, K>,
1730}
1731
1732/// The 4-rows-per-pass batched one-seed scalar (`wide::f64x4` lanes).
1733pub type OneSeedBatch<const K: usize> = OneSeedLane<wide::f64x4, K>;
1734
1735impl<L: Lane, const K: usize> OneSeedLane<L, K> {
1736    /// A constant: base = `constant(c)`, ε-part zero (mirrors [`OneSeed::constant`]).
1737    #[inline]
1738    pub fn constant(c: L) -> Self {
1739        OneSeedLane {
1740            base: Order2Lane::constant(c),
1741            eps: Order2Lane::constant(L::splat(0.0)),
1742        }
1743    }
1744
1745    /// The seeded variable `p_axis` at (per-lane) value `value`, no ε-direction
1746    /// (mirrors [`OneSeed::variable`]).
1747    #[inline]
1748    pub fn variable(value: L, axis: usize) -> Self {
1749        OneSeedLane {
1750            base: Order2Lane::variable(value, axis),
1751            eps: Order2Lane::constant(L::splat(0.0)),
1752        }
1753    }
1754
1755    /// Seed primary `axis` at (per-lane) value `value` with ε-direction
1756    /// component `u_axis`: base = `variable(value, axis)`, eps = `constant(u_axis)`
1757    /// (mirrors [`OneSeed::seed_direction`]). With `L = f64x4`, `value` / `u_axis`
1758    /// pack the four rows' values / directions of primary `axis`.
1759    #[inline]
1760    pub fn seed_direction(value: L, axis: usize, u_axis: L) -> Self {
1761        OneSeedLane {
1762            base: Order2Lane::variable(value, axis),
1763            eps: Order2Lane::constant(u_axis),
1764        }
1765    }
1766
1767    /// The contracted-third channel after a `seed_direction(u)` evaluation:
1768    /// `out[a][b]` lane `i` is row `i`'s `Σ_c ℓ_{abc} u_c` (the ε-part Hessian).
1769    #[inline]
1770    #[must_use]
1771    pub fn contracted_third(&self) -> [[L; K]; K] {
1772        self.eps.h
1773    }
1774
1775    /// Lane-wise `self + o` (mirrors [`OneSeed::add`]).
1776    #[inline]
1777    pub fn add(&self, o: &Self) -> Self {
1778        OneSeedLane {
1779            base: self.base.add(&o.base),
1780            eps: self.eps.add(&o.eps),
1781        }
1782    }
1783
1784    /// Lane-wise `self - o` (mirrors [`OneSeed::sub`]).
1785    #[inline]
1786    pub fn sub(&self, o: &Self) -> Self {
1787        OneSeedLane {
1788            base: self.base.sub(&o.base),
1789            eps: self.eps.sub(&o.eps),
1790        }
1791    }
1792
1793    /// Lane-wise `self · o`, ε² = 0 truncation (mirrors [`OneSeed::mul`]).
1794    #[inline]
1795    pub fn mul(&self, o: &Self) -> Self {
1796        OneSeedLane {
1797            base: self.base.mul(&o.base),
1798            eps: self.base.mul(&o.eps).add(&self.eps.mul(&o.base)),
1799        }
1800    }
1801
1802    /// Negate every part (mirrors [`OneSeed::neg`]).
1803    #[inline]
1804    pub fn neg(&self) -> Self {
1805        OneSeedLane {
1806            base: self.base.neg(),
1807            eps: self.eps.neg(),
1808        }
1809    }
1810
1811    /// Multiply every part by the plain scalar `s` (mirrors [`OneSeed::scale`]).
1812    #[inline]
1813    pub fn scale(&self, s: f64) -> Self {
1814        OneSeedLane {
1815            base: self.base.scale(s),
1816            eps: self.eps.scale(s),
1817        }
1818    }
1819
1820    /// Exact order-≤2-per-part Faà di Bruno composition `f ∘ self`, given the
1821    /// per-lane outer-derivative stack `d = [f, f′, f″, f‴, f⁗]`. Term-for-term
1822    /// identical to [`OneSeed::compose_unary`]: the base reads `d[0..=2]` and the
1823    /// ε-coefficient is `f′(base)` (reads `d[1..=3]`) times `eps`.
1824    #[inline]
1825    pub fn compose_unary(&self, d: [L; 5]) -> Self {
1826        let base = self.base.compose_unary([d[0], d[1], d[2]]);
1827        let fprime = self.base.compose_unary([d[1], d[2], d[3]]);
1828        let eps = fprime.mul(&self.eps);
1829        OneSeedLane { base, eps }
1830    }
1831
1832    /// `e^self`, per-lane stack `[e, e, e, e, e]` (matches [`JetScalar::exp`]).
1833    #[inline]
1834    pub fn exp(&self) -> Self {
1835        let d = self.base.v.unary5(|u| {
1836            let e = u.exp();
1837            [e, e, e, e, e]
1838        });
1839        self.compose_unary(d)
1840    }
1841
1842    /// `ln(self)`; caller guarantees positivity (matches [`JetScalar::ln`]).
1843    #[inline]
1844    pub fn ln(&self) -> Self {
1845        let d = self.base.v.unary5(|u| {
1846            let r = 1.0 / u;
1847            [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
1848        });
1849        self.compose_unary(d)
1850    }
1851
1852    /// `√self`; caller guarantees positivity (matches [`JetScalar::sqrt`]).
1853    #[inline]
1854    pub fn sqrt(&self) -> Self {
1855        let d = self.base.v.unary5(|u| {
1856            let s = u.sqrt();
1857            [
1858                s,
1859                0.5 / s,
1860                -0.25 / (u * s),
1861                0.375 / (u * u * s),
1862                -0.9375 / (u * u * u * s),
1863            ]
1864        });
1865        self.compose_unary(d)
1866    }
1867
1868    /// `1/self` (matches [`JetScalar::recip`]).
1869    #[inline]
1870    pub fn recip(&self) -> Self {
1871        let d = self.base.v.unary5(|u| {
1872            let r = 1.0 / u;
1873            let r2 = r * r;
1874            [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
1875        });
1876        self.compose_unary(d)
1877    }
1878
1879    /// `self^a` for real `a`; caller guarantees a positive base (matches
1880    /// [`JetScalar::powf`]).
1881    #[inline]
1882    pub fn powf(&self, a: f64) -> Self {
1883        let d = self.base.v.unary5(|u| {
1884            [
1885                u.powf(a),
1886                a * u.powf(a - 1.0),
1887                a * (a - 1.0) * u.powf(a - 2.0),
1888                a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
1889                a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
1890            ]
1891        });
1892        self.compose_unary(d)
1893    }
1894
1895    /// `ln Γ(self)`; caller guarantees positivity (matches [`JetScalar::ln_gamma`],
1896    /// same hand-certified stack).
1897    #[inline]
1898    pub fn ln_gamma(&self) -> Self {
1899        let d = self
1900            .base
1901            .v
1902            .unary5(crate::jet_tower::ln_gamma_derivative_stack);
1903        self.compose_unary(d)
1904    }
1905
1906    /// `ψ(self)` digamma; caller guarantees positivity (matches
1907    /// [`JetScalar::digamma`], same hand-certified stack).
1908    #[inline]
1909    pub fn digamma(&self) -> Self {
1910        let d = self
1911            .base
1912            .v
1913            .unary5(crate::jet_tower::digamma_derivative_stack);
1914        self.compose_unary(d)
1915    }
1916}
1917
1918impl<const K: usize> OneSeedBatch<K> {
1919    /// Extract lane `i`'s parts as a production [`OneSeed<K>`]. Lane `i` is
1920    /// `to_bits`-identical to evaluating the same program at [`OneSeed<K>`] on
1921    /// row `i` (see `batch_tests`).
1922    #[inline]
1923    #[must_use]
1924    pub fn lane(&self, i: usize) -> OneSeed<K> {
1925        OneSeed {
1926            base: self.base.lane(i),
1927            eps: self.eps.lane(i),
1928        }
1929    }
1930}
1931
1932// ── TwoSeed<K>: two-seed, contracted fourth (doc §A.3) ──────────────────
1933
1934/// Two-seed scalar: an [`Order2`] base plus TWO nilpotents ε, δ
1935/// (`ε² = δ² = 0`, `εδ` retained) — four [`Order2`] parts
1936/// `s = base + ε·eps + δ·del + εδ·eps_del`.
1937///
1938/// Product truncates `ε² = δ² = 0` (doc §A.3): each part is built from
1939/// [`Order2`] products of the four input parts. Composition picks up
1940/// successively higher outer derivatives, the cross part carrying the second
1941/// Faà di Bruno term `f''·eps·del + f'·eps_del`.
1942///
1943/// Seed each primary with [`seed`](Self::seed): base = `variable(x, axis)`,
1944/// eps = `constant(u_axis)`, del = `constant(v_axis)`, eps_del = `constant(0)`.
1945/// Then the εδ-component of the evaluated Hessian channel is the contracted
1946/// fourth `[eps_del.h][a][b] = Σ_{cd} ℓ_{abcd} u_c v_d` — exactly
1947/// `row_fourth_contracted(u, v)`, without materialising `t4`.
1948#[derive(Clone, Copy, Debug)]
1949pub struct TwoSeed<const K: usize> {
1950    /// The `ε⁰δ⁰` part: value / grad / Hessian of `ℓ`.
1951    pub base: Order2<K>,
1952    /// The `ε¹δ⁰` part.
1953    pub eps: Order2<K>,
1954    /// The `ε⁰δ¹` part.
1955    pub del: Order2<K>,
1956    /// The `ε¹δ¹` part. After a `seed(u, v)` evaluation,
1957    /// `eps_del.h[a][b] = Σ_{cd} ℓ_{abcd} u_c v_d`.
1958    pub eps_del: Order2<K>,
1959}
1960
1961impl<const K: usize> TwoSeed<K> {
1962    /// Seed primary `axis` at value `x` with ε-direction `u_axis` and
1963    /// δ-direction `v_axis`:
1964    /// `p_axis = p_axis⁰ + x-seed + ε·u_axis + δ·v_axis` (doc §A.3 "Seeding").
1965    pub fn seed(x: f64, axis: usize, u_axis: f64, v_axis: f64) -> Self {
1966        TwoSeed {
1967            base: Order2::variable(x, axis),
1968            eps: Order2::constant(u_axis),
1969            del: Order2::constant(v_axis),
1970            eps_del: Order2::constant(0.0),
1971        }
1972    }
1973
1974    /// The contracted-fourth channel after a `seed(u, v)` evaluation:
1975    /// `out[a][b] = Σ_{cd} ℓ_{abcd} u_c v_d`, i.e. the εδ-coefficient's Hessian.
1976    pub fn contracted_fourth(&self) -> [[f64; K]; K] {
1977        self.eps_del.h()
1978    }
1979}
1980
1981impl<const K: usize> JetScalar<K> for TwoSeed<K> {
1982    fn constant(c: f64) -> Self {
1983        TwoSeed {
1984            base: Order2::constant(c),
1985            eps: Order2::constant(0.0),
1986            del: Order2::constant(0.0),
1987            eps_del: Order2::constant(0.0),
1988        }
1989    }
1990    fn variable(x: f64, axis: usize) -> Self {
1991        TwoSeed {
1992            base: Order2::variable(x, axis),
1993            eps: Order2::constant(0.0),
1994            del: Order2::constant(0.0),
1995            eps_del: Order2::constant(0.0),
1996        }
1997    }
1998    fn value(&self) -> f64 {
1999        self.base.value()
2000    }
2001    fn add(&self, o: &Self) -> Self {
2002        TwoSeed {
2003            base: self.base.add(&o.base),
2004            eps: self.eps.add(&o.eps),
2005            del: self.del.add(&o.del),
2006            eps_del: self.eps_del.add(&o.eps_del),
2007        }
2008    }
2009    fn sub(&self, o: &Self) -> Self {
2010        TwoSeed {
2011            base: self.base.sub(&o.base),
2012            eps: self.eps.sub(&o.eps),
2013            del: self.del.sub(&o.del),
2014            eps_del: self.eps_del.sub(&o.eps_del),
2015        }
2016    }
2017    fn mul(&self, o: &Self) -> Self {
2018        let a = self;
2019        let b = o;
2020        // Truncate ε² = δ² = 0 (doc §A.3 product table).
2021        let base = a.base.mul(&b.base);
2022        let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
2023        let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
2024        let eps_del = a
2025            .base
2026            .mul(&b.eps_del)
2027            .add(&a.eps.mul(&b.del))
2028            .add(&a.del.mul(&b.eps))
2029            .add(&a.eps_del.mul(&b.base));
2030        TwoSeed {
2031            base,
2032            eps,
2033            del,
2034            eps_del,
2035        }
2036    }
2037    fn neg(&self) -> Self {
2038        TwoSeed {
2039            base: self.base.neg(),
2040            eps: self.eps.neg(),
2041            del: self.del.neg(),
2042            eps_del: self.eps_del.neg(),
2043        }
2044    }
2045    fn scale(&self, s: f64) -> Self {
2046        TwoSeed {
2047            base: self.base.scale(s),
2048            eps: self.eps.scale(s),
2049            del: self.del.scale(s),
2050            eps_del: self.eps_del.scale(s),
2051        }
2052    }
2053    fn compose_unary(&self, d: [f64; 5]) -> Self {
2054        // f(s) with s = base + ε eps + δ del + εδ eps_del, ε²=δ²=0:
2055        //   f(s) = f(base)
2056        //        + ε · f'(base)·eps
2057        //        + δ · f'(base)·del
2058        //        + εδ · ( f''(base)·eps·del + f'(base)·eps_del ).
2059        // Each f^{(r)}(base) is the Order2 composition of base with the stack
2060        // shifted r entries (doc §A.3 composition). Order2 reads only the
2061        // leading three entries of whatever stack it is handed, so the trailing
2062        // padding slots are unused (the fixed-length array makes this total).
2063        let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
2064        let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]); // f'(base) as Order2
2065        let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]); // f''(base) as Order2
2066        let eps = fprime.mul(&self.eps);
2067        let del = fprime.mul(&self.del);
2068        let eps_del = fsecond
2069            .mul(&self.eps)
2070            .mul(&self.del)
2071            .add(&fprime.mul(&self.eps_del));
2072        TwoSeed {
2073            base,
2074            eps,
2075            del,
2076            eps_del,
2077        }
2078    }
2079}
2080
2081// ── TwoSeedLane<L, K>: lane-batched two-seed, contracted fourth (doc §A.3) ─
2082
2083/// Lane-batched [`TwoSeed`]: the same two-seed scalar with its four [`Order2`]
2084/// parts re-typed to [`Order2Lane<L, K>`], so one `L = f64x4` instance carries
2085/// FOUR rows' contracted-fourth evaluations per vector pass.
2086///
2087/// Every operation is a term-for-term structural re-type of the scalar
2088/// [`TwoSeed`] ops onto the lane-implemented [`Order2Lane`] algebra. With
2089/// `L = f64`, `TwoSeedLane<f64, K>` is `to_bits`-identical to [`TwoSeed<K>`];
2090/// with `L = f64x4`, lane `i` is `to_bits`-identical to that (see `batch_tests`).
2091#[derive(Clone, Copy, Debug)]
2092pub struct TwoSeedLane<L: Lane, const K: usize> {
2093    /// The `ε⁰δ⁰` part.
2094    pub base: Order2Lane<L, K>,
2095    /// The `ε¹δ⁰` part.
2096    pub eps: Order2Lane<L, K>,
2097    /// The `ε⁰δ¹` part.
2098    pub del: Order2Lane<L, K>,
2099    /// The `ε¹δ¹` part. After a `seed(u, v)` evaluation, `eps_del.h[a][b]`
2100    /// lane `i` is row `i`'s `Σ_{cd} ℓ_{abcd} u_c v_d`.
2101    pub eps_del: Order2Lane<L, K>,
2102}
2103
2104/// The 4-rows-per-pass batched two-seed scalar (`wide::f64x4` lanes).
2105pub type TwoSeedBatch<const K: usize> = TwoSeedLane<wide::f64x4, K>;
2106
2107impl<L: Lane, const K: usize> TwoSeedLane<L, K> {
2108    /// A constant: base = `constant(c)`, all seed parts zero (mirrors
2109    /// [`TwoSeed::constant`]).
2110    #[inline]
2111    pub fn constant(c: L) -> Self {
2112        let z = Order2Lane::constant(L::splat(0.0));
2113        TwoSeedLane {
2114            base: Order2Lane::constant(c),
2115            eps: z,
2116            del: z,
2117            eps_del: z,
2118        }
2119    }
2120
2121    /// The seeded variable `p_axis` at (per-lane) value `value`, no ε/δ direction
2122    /// (mirrors [`TwoSeed::variable`]).
2123    #[inline]
2124    pub fn variable(value: L, axis: usize) -> Self {
2125        let z = Order2Lane::constant(L::splat(0.0));
2126        TwoSeedLane {
2127            base: Order2Lane::variable(value, axis),
2128            eps: z,
2129            del: z,
2130            eps_del: z,
2131        }
2132    }
2133
2134    /// Seed primary `axis` at (per-lane) value `value` with ε-direction `u_axis`
2135    /// and δ-direction `v_axis` (mirrors [`TwoSeed::seed`]). With `L = f64x4`,
2136    /// each argument packs the four rows' values for primary `axis`.
2137    #[inline]
2138    pub fn seed(value: L, axis: usize, u_axis: L, v_axis: L) -> Self {
2139        TwoSeedLane {
2140            base: Order2Lane::variable(value, axis),
2141            eps: Order2Lane::constant(u_axis),
2142            del: Order2Lane::constant(v_axis),
2143            eps_del: Order2Lane::constant(L::splat(0.0)),
2144        }
2145    }
2146
2147    /// The contracted-fourth channel after a `seed(u, v)` evaluation:
2148    /// `out[a][b]` lane `i` is row `i`'s `Σ_{cd} ℓ_{abcd} u_c v_d`
2149    /// (the εδ-part Hessian).
2150    #[inline]
2151    #[must_use]
2152    pub fn contracted_fourth(&self) -> [[L; K]; K] {
2153        self.eps_del.h
2154    }
2155
2156    /// Lane-wise `self + o` (mirrors [`TwoSeed::add`]).
2157    #[inline]
2158    pub fn add(&self, o: &Self) -> Self {
2159        TwoSeedLane {
2160            base: self.base.add(&o.base),
2161            eps: self.eps.add(&o.eps),
2162            del: self.del.add(&o.del),
2163            eps_del: self.eps_del.add(&o.eps_del),
2164        }
2165    }
2166
2167    /// Lane-wise `self - o` (mirrors [`TwoSeed::sub`]).
2168    #[inline]
2169    pub fn sub(&self, o: &Self) -> Self {
2170        TwoSeedLane {
2171            base: self.base.sub(&o.base),
2172            eps: self.eps.sub(&o.eps),
2173            del: self.del.sub(&o.del),
2174            eps_del: self.eps_del.sub(&o.eps_del),
2175        }
2176    }
2177
2178    /// Lane-wise `self · o`, ε² = δ² = 0 truncation (mirrors [`TwoSeed::mul`]).
2179    #[inline]
2180    pub fn mul(&self, o: &Self) -> Self {
2181        let a = self;
2182        let b = o;
2183        let base = a.base.mul(&b.base);
2184        let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
2185        let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
2186        let eps_del = a
2187            .base
2188            .mul(&b.eps_del)
2189            .add(&a.eps.mul(&b.del))
2190            .add(&a.del.mul(&b.eps))
2191            .add(&a.eps_del.mul(&b.base));
2192        TwoSeedLane {
2193            base,
2194            eps,
2195            del,
2196            eps_del,
2197        }
2198    }
2199
2200    /// Negate every part (mirrors [`TwoSeed::neg`]).
2201    #[inline]
2202    pub fn neg(&self) -> Self {
2203        TwoSeedLane {
2204            base: self.base.neg(),
2205            eps: self.eps.neg(),
2206            del: self.del.neg(),
2207            eps_del: self.eps_del.neg(),
2208        }
2209    }
2210
2211    /// Multiply every part by the plain scalar `s` (mirrors [`TwoSeed::scale`]).
2212    #[inline]
2213    pub fn scale(&self, s: f64) -> Self {
2214        TwoSeedLane {
2215            base: self.base.scale(s),
2216            eps: self.eps.scale(s),
2217            del: self.del.scale(s),
2218            eps_del: self.eps_del.scale(s),
2219        }
2220    }
2221
2222    /// Exact composition `f ∘ self`, given the per-lane outer-derivative stack
2223    /// `d = [f, f′, f″, f‴, f⁗]`. Term-for-term identical to
2224    /// [`TwoSeed::compose_unary`]: base reads `d[0..=2]`, `f′(base)` reads
2225    /// `d[1..=3]`, `f″(base)` reads `d[2..=4]`, and the cross part carries
2226    /// `f″·eps·del + f′·eps_del`.
2227    #[inline]
2228    pub fn compose_unary(&self, d: [L; 5]) -> Self {
2229        let base = self.base.compose_unary([d[0], d[1], d[2]]);
2230        let fprime = self.base.compose_unary([d[1], d[2], d[3]]);
2231        let fsecond = self.base.compose_unary([d[2], d[3], d[4]]);
2232        let eps = fprime.mul(&self.eps);
2233        let del = fprime.mul(&self.del);
2234        let eps_del = fsecond
2235            .mul(&self.eps)
2236            .mul(&self.del)
2237            .add(&fprime.mul(&self.eps_del));
2238        TwoSeedLane {
2239            base,
2240            eps,
2241            del,
2242            eps_del,
2243        }
2244    }
2245
2246    /// `e^self`, per-lane stack `[e; 5]` (matches [`JetScalar::exp`]).
2247    #[inline]
2248    pub fn exp(&self) -> Self {
2249        let d = self.base.v.unary5(|u| {
2250            let e = u.exp();
2251            [e, e, e, e, e]
2252        });
2253        self.compose_unary(d)
2254    }
2255
2256    /// `ln(self)`; caller guarantees positivity (matches [`JetScalar::ln`]).
2257    #[inline]
2258    pub fn ln(&self) -> Self {
2259        let d = self.base.v.unary5(|u| {
2260            let r = 1.0 / u;
2261            [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
2262        });
2263        self.compose_unary(d)
2264    }
2265
2266    /// `√self`; caller guarantees positivity (matches [`JetScalar::sqrt`]).
2267    #[inline]
2268    pub fn sqrt(&self) -> Self {
2269        let d = self.base.v.unary5(|u| {
2270            let s = u.sqrt();
2271            [
2272                s,
2273                0.5 / s,
2274                -0.25 / (u * s),
2275                0.375 / (u * u * s),
2276                -0.9375 / (u * u * u * s),
2277            ]
2278        });
2279        self.compose_unary(d)
2280    }
2281
2282    /// `1/self` (matches [`JetScalar::recip`]).
2283    #[inline]
2284    pub fn recip(&self) -> Self {
2285        let d = self.base.v.unary5(|u| {
2286            let r = 1.0 / u;
2287            let r2 = r * r;
2288            [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
2289        });
2290        self.compose_unary(d)
2291    }
2292
2293    /// `self^a` for real `a`; caller guarantees a positive base (matches
2294    /// [`JetScalar::powf`]).
2295    #[inline]
2296    pub fn powf(&self, a: f64) -> Self {
2297        let d = self.base.v.unary5(|u| {
2298            [
2299                u.powf(a),
2300                a * u.powf(a - 1.0),
2301                a * (a - 1.0) * u.powf(a - 2.0),
2302                a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
2303                a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
2304            ]
2305        });
2306        self.compose_unary(d)
2307    }
2308
2309    /// `ln Γ(self)`; caller guarantees positivity (matches [`JetScalar::ln_gamma`]).
2310    #[inline]
2311    pub fn ln_gamma(&self) -> Self {
2312        let d = self
2313            .base
2314            .v
2315            .unary5(crate::jet_tower::ln_gamma_derivative_stack);
2316        self.compose_unary(d)
2317    }
2318
2319    /// `ψ(self)` digamma; caller guarantees positivity (matches
2320    /// [`JetScalar::digamma`]).
2321    #[inline]
2322    pub fn digamma(&self) -> Self {
2323        let d = self
2324            .base
2325            .v
2326            .unary5(crate::jet_tower::digamma_derivative_stack);
2327        self.compose_unary(d)
2328    }
2329}
2330
2331impl<const K: usize> TwoSeedBatch<K> {
2332    /// Extract lane `i`'s parts as a production [`TwoSeed<K>`]. Lane `i` is
2333    /// `to_bits`-identical to evaluating the same program at [`TwoSeed<K>`] on
2334    /// row `i` (see `batch_tests`).
2335    #[inline]
2336    #[must_use]
2337    pub fn lane(&self, i: usize) -> TwoSeed<K> {
2338        TwoSeed {
2339            base: self.base.lane(i),
2340            eps: self.eps.lane(i),
2341            del: self.del.lane(i),
2342            eps_del: self.eps_del.lane(i),
2343        }
2344    }
2345}
2346
2347// ── Tower3<K>: value / gradient / Hessian / third tensor ────────────────
2348
2349/// The order-≤3 [`crate::jet_tower::Tower3`] is also a [`JetScalar`]. It serves
2350/// consumers that read `.t3` but never `.t4`, avoiding the fourth-tensor
2351/// product/composition work while preserving the lower channels
2352/// bit-for-bit against [`crate::jet_tower::Tower4`].
2353impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower3<K> {
2354    fn constant(c: f64) -> Self {
2355        crate::jet_tower::Tower3::constant(c)
2356    }
2357    fn variable(x: f64, axis: usize) -> Self {
2358        crate::jet_tower::Tower3::variable(x, axis)
2359    }
2360    fn value(&self) -> f64 {
2361        self.v
2362    }
2363    fn add(&self, o: &Self) -> Self {
2364        *self + *o
2365    }
2366    fn sub(&self, o: &Self) -> Self {
2367        *self + o.scale(-1.0)
2368    }
2369    fn mul(&self, o: &Self) -> Self {
2370        crate::jet_tower::Tower3::mul(self, o)
2371    }
2372    fn neg(&self) -> Self {
2373        self.scale(-1.0)
2374    }
2375    fn scale(&self, s: f64) -> Self {
2376        crate::jet_tower::Tower3::scale(self, s)
2377    }
2378    fn compose_unary(&self, d: [f64; 5]) -> Self {
2379        crate::jet_tower::Tower3::compose_unary(self, [d[0], d[1], d[2], d[3]])
2380    }
2381}
2382
2383// ── Tower4<K>: full dense tower as a JetScalar (the all-channels scalar) ─
2384
2385/// The full dense [`crate::jet_tower::Tower4`] is itself a [`JetScalar`]: it
2386/// carries EVERY channel, so a row expression written ONCE against [`JetScalar`]
2387/// can be evaluated at `Tower4` to obtain the full `(v, g, H, t3, t4)` in one
2388/// pass. This is BOTH the #932 oracle ground truth the packed [`Order2`] /
2389/// [`OneSeed`] / [`TwoSeed`] scalars are pinned against, AND a production scalar:
2390/// a family whose uncontracted third / fourth derivative tensors are needed
2391/// (the BMS rigid `third_full` / `fourth_full` caches) evaluates the SAME
2392/// generic row-NLL expression at `Tower4` and reads `.t3` / `.t4` off the
2393/// result — so the dense tensors come from the single source of truth, not a
2394/// separately hand-written jet. The packed scalars serve the consumers that
2395/// need only `(v, g, H)` (`Order2`) or one / two contractions
2396/// (`OneSeed` / `TwoSeed`) without paying for the dense tensors.
2397impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower4<K> {
2398    fn constant(c: f64) -> Self {
2399        crate::jet_tower::Tower4::constant(c)
2400    }
2401    fn variable(x: f64, axis: usize) -> Self {
2402        crate::jet_tower::Tower4::variable(x, axis)
2403    }
2404    fn value(&self) -> f64 {
2405        self.v
2406    }
2407    fn add(&self, o: &Self) -> Self {
2408        *self + *o
2409    }
2410    fn sub(&self, o: &Self) -> Self {
2411        *self - *o
2412    }
2413    fn mul(&self, o: &Self) -> Self {
2414        crate::jet_tower::Tower4::mul(self, o)
2415    }
2416    fn neg(&self) -> Self {
2417        self.scale(-1.0)
2418    }
2419    fn scale(&self, s: f64) -> Self {
2420        crate::jet_tower::Tower4::scale(self, s)
2421    }
2422    fn compose_unary(&self, d: [f64; 5]) -> Self {
2423        crate::jet_tower::Tower4::compose_unary(self, d)
2424    }
2425}
2426
2427#[cfg(test)]
2428mod tests {
2429    use super::*;
2430    use crate::jet_tower::{RowNllProgram, Tower4, evaluate_program};
2431
2432    /// A small polynomial-plus-unary row expression written ONCE, generically
2433    /// over `S: JetScalar<2>`, so it can be evaluated against every scalar:
2434    /// `ℓ = (e^{p0·p1} + 2) · √(p0·p0 + 1) − p1·p1·0.5`.
2435    /// Exercises mul, add/sub, scale, exp, sqrt — every algebra op.
2436    fn row_expr<S: JetScalar<2>>(p: &[S; 2]) -> S {
2437        let g = p[0].mul(&p[1]).exp();
2438        let inner = g.add(&S::constant(2.0));
2439        let radic = p[0].mul(&p[0]).add(&S::constant(1.0)).sqrt();
2440        inner.mul(&radic).sub(&p[1].mul(&p[1]).scale(0.5))
2441    }
2442
2443    /// The same expression as a Tower4 `RowNllProgram`, the ground-truth tower.
2444    struct ExprProgram {
2445        p: [f64; 2],
2446    }
2447    impl RowNllProgram<2> for ExprProgram {
2448        fn n_rows(&self) -> usize {
2449            1
2450        }
2451        fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2452            if row >= self.n_rows() {
2453                return Err(format!("ExprProgram: row {row} out of range"));
2454            }
2455            Ok(self.p)
2456        }
2457        fn row_nll(&self, row: usize, p: &[Tower4<2>; 2]) -> Result<Tower4<2>, String> {
2458            if row >= self.n_rows() {
2459                return Err(format!("ExprProgram: row {row} out of range"));
2460            }
2461            Ok(row_expr(p))
2462        }
2463    }
2464
2465    const SEED: [f64; 2] = [0.37, -0.81];
2466    const U: [f64; 2] = [0.6, -0.2];
2467    const V: [f64; 2] = [-0.4, 1.1];
2468    const TOL: f64 = 1e-10;
2469
2470    fn close(a: f64, b: f64, label: &str) {
2471        let band = TOL + TOL * a.abs().max(b.abs());
2472        assert!(
2473            (a - b).abs() <= band,
2474            "{label}: {a:+.15e} vs {b:+.15e} (band {band:.3e})"
2475        );
2476    }
2477
2478    fn tower() -> Tower4<2> {
2479        evaluate_program(&ExprProgram { p: SEED }, 0).expect("tower")
2480    }
2481
2482    /// Order2 reproduces Tower4's value/grad/Hessian channels exactly.
2483    #[test]
2484    fn order2_matches_tower_value_grad_hessian() {
2485        let t = tower();
2486        let vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
2487        let s = row_expr(&vars);
2488        close(s.value(), t.v, "value");
2489        for a in 0..2 {
2490            close(s.0.g[a], t.g[a], &format!("grad[{a}]"));
2491            for b in 0..2 {
2492                close(s.h()[a][b], t.h[a][b], &format!("hess[{a}][{b}]"));
2493            }
2494        }
2495    }
2496
2497    /// The `compose_unary_with` seam on a scalar jet is `to_bits`-identical to
2498    /// the explicit `compose_unary(stack_fn(value))` — the contract the batch
2499    /// arm (`Tower{3,4}Lane::compose_unary_with`) lane-matches. Exercised on
2500    /// [`Order2`] across `K ∈ {2,3,4,9}`, ≥ 4000 random seeded inputs.
2501    #[test]
2502    fn compose_unary_with_scalar_seam_bit_identical() {
2503        fn rand_unit(state: &mut u64) -> f64 {
2504            let mut x = *state;
2505            x ^= x << 13;
2506            x ^= x >> 7;
2507            x ^= x << 17;
2508            *state = x;
2509            2.0 * ((x >> 11) as f64 / ((1u64 << 53) as f64)) - 1.0
2510        }
2511        // A base-value-dependent finite stack standing in for a family stack.
2512        fn stack(u: f64) -> [f64; 5] {
2513            [
2514                u.sin(),
2515                u.cos(),
2516                (2.0 * u).sin(),
2517                (0.5 * u).cos(),
2518                u * u - 0.3,
2519            ]
2520        }
2521        fn run<const K: usize>(state: &mut u64, n: usize) -> usize {
2522            for _ in 0..n {
2523                // A non-trivial Order2<K> jet: a seeded variable pushed through a
2524                // couple of algebra ops so g/h are dense, then exercise the seam.
2525                let base = rand_unit(state);
2526                let mut s = Order2::<K>::variable(base, 0);
2527                for a in 1..K {
2528                    s = JetScalar::mul(&s, &Order2::<K>::variable(rand_unit(state), a));
2529                }
2530                let with = s.compose_unary_with(stack);
2531                let explicit = s.compose_unary(stack(s.value()));
2532                assert_eq!(with.value().to_bits(), explicit.value().to_bits(), "value");
2533                for a in 0..K {
2534                    assert_eq!(with.g()[a].to_bits(), explicit.g()[a].to_bits(), "g[{a}]");
2535                    for b in 0..K {
2536                        assert_eq!(
2537                            with.h()[a][b].to_bits(),
2538                            explicit.h()[a][b].to_bits(),
2539                            "h[{a}][{b}]"
2540                        );
2541                    }
2542                }
2543            }
2544            n
2545        }
2546        let mut st = 0x9e37_79b9_7f4a_7c15u64;
2547        let total = run::<2>(&mut st, 1100)
2548            + run::<3>(&mut st, 1100)
2549            + run::<4>(&mut st, 1100)
2550            + run::<9>(&mut st, 1100);
2551        assert_eq!(total, 4400);
2552    }
2553
2554    /// OneSeed's ε-Hessian is the contracted third Σ_c ℓ_{abc} u_c, matching
2555    /// `Tower4::third_contracted(u)`. Base channels also match the tower.
2556    #[test]
2557    fn one_seed_matches_tower_third_contracted() {
2558        let t = tower();
2559        let truth = t.third_contracted(&U);
2560        let vars: [OneSeed<2>; 2] =
2561            std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
2562        let s = row_expr(&vars);
2563        // Base channels are the plain (v, g, H).
2564        close(s.value(), t.v, "value");
2565        for a in 0..2 {
2566            for b in 0..2 {
2567                close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
2568            }
2569        }
2570        let third = s.contracted_third();
2571        for a in 0..2 {
2572            for b in 0..2 {
2573                close(third[a][b], truth[a][b], &format!("third[{a}][{b}]"));
2574            }
2575        }
2576    }
2577
2578    /// TwoSeed's εδ-Hessian is the contracted fourth Σ_{cd} ℓ_{abcd} u_c v_d,
2579    /// matching `Tower4::fourth_contracted(u, v)`. The ε / δ single-seed parts
2580    /// reproduce the two third contractions Σ_c ℓ_{abc} u_c and …v_d.
2581    #[test]
2582    fn two_seed_matches_tower_fourth_contracted() {
2583        let t = tower();
2584        let truth4 = t.fourth_contracted(&U, &V);
2585        let truth3_u = t.third_contracted(&U);
2586        let truth3_v = t.third_contracted(&V);
2587        let vars: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
2588        let s = row_expr(&vars);
2589        close(s.value(), t.v, "value");
2590        for a in 0..2 {
2591            close(s.base.0.g[a], t.g[a], &format!("grad[{a}]"));
2592            for b in 0..2 {
2593                close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
2594                close(
2595                    s.eps.h()[a][b],
2596                    truth3_u[a][b],
2597                    &format!("eps third_u[{a}][{b}]"),
2598                );
2599                close(
2600                    s.del.h()[a][b],
2601                    truth3_v[a][b],
2602                    &format!("del third_v[{a}][{b}]"),
2603                );
2604            }
2605        }
2606        let fourth = s.contracted_fourth();
2607        for a in 0..2 {
2608            for b in 0..2 {
2609                close(fourth[a][b], truth4[a][b], &format!("fourth[{a}][{b}]"));
2610            }
2611        }
2612    }
2613
2614    /// The generic `row_nll_generic` seam (added to Tower4's program trait
2615    /// surface) evaluates the SAME expression on each scalar and extracts the
2616    /// channel a consumer asks for, agreeing with the direct Tower4 contraction.
2617    #[test]
2618    fn generic_program_seam_matches_tower_for_every_channel() {
2619        let t = tower();
2620        // Order2 via generic seam.
2621        let o2: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
2622        let so2 = row_expr(&o2);
2623        close(so2.value(), t.v, "seam order2 value");
2624        // OneSeed third.
2625        let os: [OneSeed<2>; 2] =
2626            std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
2627        let third = row_expr(&os).contracted_third();
2628        let truth3 = t.third_contracted(&U);
2629        for a in 0..2 {
2630            for b in 0..2 {
2631                close(third[a][b], truth3[a][b], &format!("seam third[{a}][{b}]"));
2632            }
2633        }
2634        // TwoSeed fourth.
2635        let ts: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
2636        let fourth = row_expr(&ts).contracted_fourth();
2637        let truth4 = t.fourth_contracted(&U, &V);
2638        for a in 0..2 {
2639            for b in 0..2 {
2640                close(
2641                    fourth[a][b],
2642                    truth4[a][b],
2643                    &format!("seam fourth[{a}][{b}]"),
2644                );
2645            }
2646        }
2647    }
2648
2649    /// The (test-only) `Tower4: JetScalar` impl is the all-channels oracle scalar:
2650    /// evaluating the SAME generic `row_expr` at `S = Tower4` (through the
2651    /// `JetScalar` trait ops) must reproduce, channel-for-channel, the `Tower4`
2652    /// obtained from the `RowNllProgram` / inherent-operator path
2653    /// (`evaluate_program`). This pins that the trait impl delegates faithfully to
2654    /// the inherent `Tower4` arithmetic (so the contracted-scalar oracles above,
2655    /// which compare against `evaluate_program`'s tower, are comparing against the
2656    /// same algebra the `JetScalar` interface exposes).
2657    #[test]
2658    fn tower4_as_jetscalar_matches_program_tower_all_channels() {
2659        let t = tower();
2660        let vars: [Tower4<2>; 2] = std::array::from_fn(|a| Tower4::variable(SEED[a], a));
2661        let s = row_expr(&vars);
2662        close(s.v, t.v, "tower-jetscalar value");
2663        for a in 0..2 {
2664            close(s.g[a], t.g[a], &format!("tower-jetscalar grad[{a}]"));
2665            for b in 0..2 {
2666                close(
2667                    s.h[a][b],
2668                    t.h[a][b],
2669                    &format!("tower-jetscalar hess[{a}][{b}]"),
2670                );
2671                for c in 0..2 {
2672                    close(
2673                        s.t3[a][b][c],
2674                        t.t3[a][b][c],
2675                        &format!("tower-jetscalar t3[{a}][{b}][{c}]"),
2676                    );
2677                    for d in 0..2 {
2678                        close(
2679                            s.t4[a][b][c][d],
2680                            t.t4[a][b][c][d],
2681                            &format!("tower-jetscalar t4[{a}][{b}][{c}][{d}]"),
2682                        );
2683                    }
2684                }
2685            }
2686        }
2687    }
2688
2689    /// Arena-backed runtime directional jets must reproduce the fixed packed
2690    /// algebras channel-for-channel. This isolates the runtime scalar algebra
2691    /// from every family row program before the SLS/SAE integration oracles.
2692    #[test]
2693    fn runtime_directional_jets_match_fixed_packed_algebra_932() {
2694        fn expression<'arena, S: RuntimeJetScalar<'arena>>(vars: &[S]) -> S {
2695            let bilinear = vars[0].mul(&vars[1]);
2696            let curved = vars[2].scale(0.7).add(&vars[3].mul(&vars[3]).scale(-0.2));
2697            bilinear
2698                .add(&curved)
2699                .exp()
2700                .mul(&vars[4].compose_unary([0.4, -0.3, 0.2, -0.1, 0.05]))
2701        }
2702
2703        const K: usize = 5;
2704        let values = [0.2, -0.7, 0.4, 1.1, -0.3];
2705        let direction_u = [0.5, -0.2, 0.7, -0.4, 0.1];
2706        let direction_v = [-0.3, 0.8, 0.2, 0.6, -0.5];
2707        let close = |actual: f64, expected: f64| {
2708            let tolerance = 1.0e-13 * (1.0 + actual.abs().max(expected.abs()));
2709            assert!((actual - expected).abs() <= tolerance);
2710        };
2711
2712        let fixed_one: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
2713            .map(|axis| FixedRuntimeJet {
2714                inner: OneSeed::seed_direction(values[axis], axis, direction_u[axis]),
2715            })
2716            .collect();
2717        let arena_one = DynamicJetArena::new();
2718        let dynamic_one: Vec<DynamicOneSeed<'_>> = (0..K)
2719            .map(|axis| {
2720                DynamicOneSeed::seed_direction(values[axis], axis, direction_u[axis], K, &arena_one)
2721            })
2722            .collect();
2723        let fixed_third = expression(&fixed_one).into_inner().contracted_third();
2724        let dynamic_third = expression(&dynamic_one);
2725        for a in 0..K {
2726            for b in 0..K {
2727                close(
2728                    dynamic_third.contracted_third()[a * K + b],
2729                    fixed_third[a][b],
2730                );
2731            }
2732        }
2733
2734        let fixed_two: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
2735            .map(|axis| FixedRuntimeJet {
2736                inner: TwoSeed::seed(values[axis], axis, direction_u[axis], direction_v[axis]),
2737            })
2738            .collect();
2739        let arena_two = DynamicJetArena::new();
2740        let dynamic_two: Vec<DynamicTwoSeed<'_>> = (0..K)
2741            .map(|axis| {
2742                DynamicTwoSeed::seed(
2743                    values[axis],
2744                    axis,
2745                    direction_u[axis],
2746                    direction_v[axis],
2747                    K,
2748                    &arena_two,
2749                )
2750            })
2751            .collect();
2752        let fixed_fourth = expression(&fixed_two).into_inner().contracted_fourth();
2753        let dynamic_fourth = expression(&dynamic_two);
2754        for a in 0..K {
2755            for b in 0..K {
2756                close(
2757                    dynamic_fourth.contracted_fourth()[a * K + b],
2758                    fixed_fourth[a][b],
2759                );
2760            }
2761        }
2762    }
2763}
2764
2765#[cfg(test)]
2766mod batch_tests {
2767    //! SIMD row-batching oracle: prove [`Order2Batch<K>`] (4 rows in
2768    //! `wide::f64x4` lanes) is `to_bits`-identical, on every value/gradient/
2769    //! Hessian channel, to the production [`Order2<K>`] evaluated per row — and
2770    //! that the new scalar field [`Order2Lane<f64, K>`] is too. Composing the two
2771    //! claims, batch lane `i` reproduces the production scalar for row `i` bit
2772    //! for bit, so the 4× throughput is a free lunch (no result change).
2773
2774    use super::{
2775        JetScalar, Lane, OneSeed, OneSeedBatch, OneSeedLane, Order2, Order2Batch, Order2Lane,
2776        TwoSeed, TwoSeedBatch, TwoSeedLane,
2777    };
2778
2779    /// The ops the witness row expression needs, so ONE generic body evaluates
2780    /// at the production [`Order2<K>`], the new scalar [`Order2Lane<f64, K>`],
2781    /// and the batched [`Order2Batch<K>`].
2782    trait RowAlg<const K: usize>: Copy {
2783        fn constant(c: f64) -> Self;
2784        fn add(&self, o: &Self) -> Self;
2785        fn sub(&self, o: &Self) -> Self;
2786        fn mul(&self, o: &Self) -> Self;
2787        fn scale(&self, s: f64) -> Self;
2788        fn exp(&self) -> Self;
2789        fn sqrt(&self) -> Self;
2790        fn recip(&self) -> Self;
2791    }
2792
2793    impl<const K: usize> RowAlg<K> for Order2<K> {
2794        fn constant(c: f64) -> Self {
2795            <Self as JetScalar<K>>::constant(c)
2796        }
2797        fn add(&self, o: &Self) -> Self {
2798            JetScalar::add(self, o)
2799        }
2800        fn sub(&self, o: &Self) -> Self {
2801            JetScalar::sub(self, o)
2802        }
2803        fn mul(&self, o: &Self) -> Self {
2804            JetScalar::mul(self, o)
2805        }
2806        fn scale(&self, s: f64) -> Self {
2807            JetScalar::scale(self, s)
2808        }
2809        fn exp(&self) -> Self {
2810            JetScalar::exp(self)
2811        }
2812        fn sqrt(&self) -> Self {
2813            JetScalar::sqrt(self)
2814        }
2815        fn recip(&self) -> Self {
2816            JetScalar::recip(self)
2817        }
2818    }
2819
2820    impl<L: Lane, const K: usize> RowAlg<K> for Order2Lane<L, K> {
2821        fn constant(c: f64) -> Self {
2822            Order2Lane::constant(L::splat(c))
2823        }
2824        fn add(&self, o: &Self) -> Self {
2825            Order2Lane::add(self, o)
2826        }
2827        fn sub(&self, o: &Self) -> Self {
2828            Order2Lane::sub(self, o)
2829        }
2830        fn mul(&self, o: &Self) -> Self {
2831            Order2Lane::mul(self, o)
2832        }
2833        fn scale(&self, s: f64) -> Self {
2834            Order2Lane::scale(self, s)
2835        }
2836        fn exp(&self) -> Self {
2837            Order2Lane::exp(self)
2838        }
2839        fn sqrt(&self) -> Self {
2840            Order2Lane::sqrt(self)
2841        }
2842        fn recip(&self) -> Self {
2843            Order2Lane::recip(self)
2844        }
2845    }
2846
2847    /// A dense witness row expression touching every algebra op (mul, add, sub,
2848    /// scale, exp, sqrt, recip) over ALL `K` primaries, so the gradient and the
2849    /// full `K×K` Hessian are dense (no trivially-zero channel). All transcend.
2850    /// arguments are kept finite/positive: `sqrt(s²+1) > 0`, `recip(exp+2) > 0`.
2851    fn row_expr<const K: usize, A: RowAlg<K>>(p: &[A; K]) -> A {
2852        let mut s = A::constant(0.3);
2853        for a in 0..K {
2854            let b = (a + 1) % K;
2855            s = s.add(&p[a].mul(&p[b]).scale(0.1 + 0.05 * a as f64));
2856        }
2857        let e = s.exp();
2858        let r = s.mul(&s).add(&A::constant(1.0)).sqrt();
2859        let denom = e.add(&A::constant(2.0));
2860        e.mul(&r).sub(&s.scale(0.5)).mul(&denom.recip())
2861    }
2862
2863    /// xorshift64 → `f64` in `[-1, 1)`.
2864    fn rand_unit(state: &mut u64) -> f64 {
2865        let mut x = *state;
2866        x ^= x << 13;
2867        x ^= x >> 7;
2868        x ^= x << 17;
2869        *state = x;
2870        let u = (x >> 11) as f64 / ((1u64 << 53) as f64); // [0, 1)
2871        2.0 * u - 1.0
2872    }
2873
2874    /// Returns the number of (batch, row) pairs whose every channel was
2875    /// verified bit-identical, so the caller can assert the expected total ran.
2876    fn check_k<const K: usize>(state: &mut u64, batches: usize) -> usize {
2877        let mut verified_rows = 0usize;
2878        for _ in 0..batches {
2879            // Four independent rows of K primary values.
2880            let rows: [[f64; K]; 4] =
2881                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
2882
2883            // Production ground truth, evaluated per row at Order2<K>.
2884            let prod: [Order2<K>; 4] = std::array::from_fn(|r| {
2885                let p: [Order2<K>; K] = std::array::from_fn(|a| Order2::variable(rows[r][a], a));
2886                row_expr(&p)
2887            });
2888
2889            // New scalar field (Order2Lane<f64>), per row.
2890            let scal: [Order2Lane<f64, K>; 4] = std::array::from_fn(|r| {
2891                let p: [Order2Lane<f64, K>; K] =
2892                    std::array::from_fn(|a| Order2Lane::variable(rows[r][a], a));
2893                row_expr(&p)
2894            });
2895
2896            // Batched: 4 rows packed into f64x4 lanes, ONE vector pass.
2897            let pbatch: [Order2Batch<K>; K] = std::array::from_fn(|a| {
2898                let packed = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
2899                Order2Batch::variable(packed, a)
2900            });
2901            let batch = row_expr(&pbatch);
2902
2903            for r in 0..4 {
2904                let g = prod[r].0;
2905                // Order2Lane<f64> == Order2<K> (bit-identical scalar field).
2906                assert_eq!(scal[r].v.to_bits(), g.v.to_bits(), "K={K} scalar v");
2907                // Batch lane r == Order2<K> for row r.
2908                let lr = batch.lane(r).0;
2909                assert_eq!(lr.v.to_bits(), g.v.to_bits(), "K={K} batch lane {r} v");
2910                for a in 0..K {
2911                    assert_eq!(
2912                        scal[r].g[a].to_bits(),
2913                        g.g[a].to_bits(),
2914                        "K={K} scalar g[{a}]"
2915                    );
2916                    assert_eq!(
2917                        lr.g[a].to_bits(),
2918                        g.g[a].to_bits(),
2919                        "K={K} batch lane {r} g[{a}]"
2920                    );
2921                    for b in 0..K {
2922                        assert_eq!(
2923                            scal[r].h[a][b].to_bits(),
2924                            g.h[a][b].to_bits(),
2925                            "K={K} scalar h[{a}][{b}]"
2926                        );
2927                        assert_eq!(
2928                            lr.h[a][b].to_bits(),
2929                            g.h[a][b].to_bits(),
2930                            "K={K} batch lane {r} h[{a}][{b}]"
2931                        );
2932                    }
2933                }
2934                verified_rows += 1;
2935            }
2936        }
2937        verified_rows
2938    }
2939
2940    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: every channel of
2941    /// every lane is `to_bits`-identical to the production scalar per row.
2942    #[test]
2943    fn batch_lanes_bit_identical_to_scalar_per_row() {
2944        let mut state = 0x9E37_79B9_7F4A_7C15_u64;
2945        let mut verified = 0usize;
2946        verified += check_k::<2>(&mut state, 2000);
2947        verified += check_k::<3>(&mut state, 2000);
2948        verified += check_k::<4>(&mut state, 2000);
2949        verified += check_k::<9>(&mut state, 2000);
2950        // 4 K-values × 2000 batches × 4 packed rows each, all bit-identical.
2951        assert_eq!(verified, 4 * 2000 * 4, "every batch row must be verified");
2952    }
2953
2954    // ── One-/two-seed lane oracles ──────────────────────────────────────────
2955    //
2956    // The same dense `row_expr` witness program runs over the SEEDED directional
2957    // scalars: the scalar `OneSeed`/`TwoSeed` per row, the `f64`-lane re-type
2958    // (`*SeedLane<f64>`), and the 4-rows-per-pass batch (`*SeedBatch`). The
2959    // headline claim is that the contracted-third / contracted-fourth channel of
2960    // every lane is `to_bits`-identical to the production scalar's per row.
2961
2962    impl<const K: usize> RowAlg<K> for OneSeed<K> {
2963        fn constant(c: f64) -> Self {
2964            <Self as JetScalar<K>>::constant(c)
2965        }
2966        fn add(&self, o: &Self) -> Self {
2967            JetScalar::add(self, o)
2968        }
2969        fn sub(&self, o: &Self) -> Self {
2970            JetScalar::sub(self, o)
2971        }
2972        fn mul(&self, o: &Self) -> Self {
2973            JetScalar::mul(self, o)
2974        }
2975        fn scale(&self, s: f64) -> Self {
2976            JetScalar::scale(self, s)
2977        }
2978        fn exp(&self) -> Self {
2979            JetScalar::exp(self)
2980        }
2981        fn sqrt(&self) -> Self {
2982            JetScalar::sqrt(self)
2983        }
2984        fn recip(&self) -> Self {
2985            JetScalar::recip(self)
2986        }
2987    }
2988
2989    impl<L: Lane, const K: usize> RowAlg<K> for OneSeedLane<L, K> {
2990        fn constant(c: f64) -> Self {
2991            OneSeedLane::constant(L::splat(c))
2992        }
2993        fn add(&self, o: &Self) -> Self {
2994            OneSeedLane::add(self, o)
2995        }
2996        fn sub(&self, o: &Self) -> Self {
2997            OneSeedLane::sub(self, o)
2998        }
2999        fn mul(&self, o: &Self) -> Self {
3000            OneSeedLane::mul(self, o)
3001        }
3002        fn scale(&self, s: f64) -> Self {
3003            OneSeedLane::scale(self, s)
3004        }
3005        fn exp(&self) -> Self {
3006            OneSeedLane::exp(self)
3007        }
3008        fn sqrt(&self) -> Self {
3009            OneSeedLane::sqrt(self)
3010        }
3011        fn recip(&self) -> Self {
3012            OneSeedLane::recip(self)
3013        }
3014    }
3015
3016    impl<const K: usize> RowAlg<K> for TwoSeed<K> {
3017        fn constant(c: f64) -> Self {
3018            <Self as JetScalar<K>>::constant(c)
3019        }
3020        fn add(&self, o: &Self) -> Self {
3021            JetScalar::add(self, o)
3022        }
3023        fn sub(&self, o: &Self) -> Self {
3024            JetScalar::sub(self, o)
3025        }
3026        fn mul(&self, o: &Self) -> Self {
3027            JetScalar::mul(self, o)
3028        }
3029        fn scale(&self, s: f64) -> Self {
3030            JetScalar::scale(self, s)
3031        }
3032        fn exp(&self) -> Self {
3033            JetScalar::exp(self)
3034        }
3035        fn sqrt(&self) -> Self {
3036            JetScalar::sqrt(self)
3037        }
3038        fn recip(&self) -> Self {
3039            JetScalar::recip(self)
3040        }
3041    }
3042
3043    impl<L: Lane, const K: usize> RowAlg<K> for TwoSeedLane<L, K> {
3044        fn constant(c: f64) -> Self {
3045            TwoSeedLane::constant(L::splat(c))
3046        }
3047        fn add(&self, o: &Self) -> Self {
3048            TwoSeedLane::add(self, o)
3049        }
3050        fn sub(&self, o: &Self) -> Self {
3051            TwoSeedLane::sub(self, o)
3052        }
3053        fn mul(&self, o: &Self) -> Self {
3054            TwoSeedLane::mul(self, o)
3055        }
3056        fn scale(&self, s: f64) -> Self {
3057            TwoSeedLane::scale(self, s)
3058        }
3059        fn exp(&self) -> Self {
3060            TwoSeedLane::exp(self)
3061        }
3062        fn sqrt(&self) -> Self {
3063            TwoSeedLane::sqrt(self)
3064        }
3065        fn recip(&self) -> Self {
3066            TwoSeedLane::recip(self)
3067        }
3068    }
3069
3070    fn check_oneseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
3071        let mut rows_checked = 0;
3072        for _ in 0..batches {
3073            let rows: [[f64; K]; 4] =
3074                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3075            // Per-row ε-direction.
3076            let u: [[f64; K]; 4] =
3077                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3078
3079            // Production ground truth (scalar OneSeed per row).
3080            let prod: [OneSeed<K>; 4] = std::array::from_fn(|r| {
3081                let p: [OneSeed<K>; K] =
3082                    std::array::from_fn(|a| OneSeed::seed_direction(rows[r][a], a, u[r][a]));
3083                row_expr(&p)
3084            });
3085
3086            // f64-lane re-type per row.
3087            let scal: [OneSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
3088                let p: [OneSeedLane<f64, K>; K] =
3089                    std::array::from_fn(|a| OneSeedLane::seed_direction(rows[r][a], a, u[r][a]));
3090                row_expr(&p)
3091            });
3092
3093            // 4-rows-per-pass batch.
3094            let pbatch: [OneSeedBatch<K>; K] = std::array::from_fn(|a| {
3095                let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
3096                let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
3097                OneSeedBatch::seed_direction(val, a, uu)
3098            });
3099            let batch = row_expr(&pbatch);
3100
3101            for r in 0..4 {
3102                let want = prod[r].contracted_third();
3103                let got_scal = scal[r].contracted_third();
3104                let got_batch = batch.lane(r).contracted_third();
3105                // Value channel too (sanity that the base program agrees).
3106                assert_eq!(
3107                    scal[r].base.v.to_bits(),
3108                    prod[r].base.value().to_bits(),
3109                    "OneSeed K={K} scalar value"
3110                );
3111                assert_eq!(
3112                    batch.lane(r).base.value().to_bits(),
3113                    prod[r].base.value().to_bits(),
3114                    "OneSeed K={K} batch lane {r} value"
3115                );
3116                for a in 0..K {
3117                    for b in 0..K {
3118                        assert_eq!(
3119                            got_scal[a][b].to_bits(),
3120                            want[a][b].to_bits(),
3121                            "OneSeed K={K} scalar third[{a}][{b}]"
3122                        );
3123                        assert_eq!(
3124                            got_batch[a][b].to_bits(),
3125                            want[a][b].to_bits(),
3126                            "OneSeed K={K} batch lane {r} third[{a}][{b}]"
3127                        );
3128                    }
3129                }
3130                rows_checked += 1;
3131            }
3132        }
3133        rows_checked
3134    }
3135
3136    fn check_twoseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
3137        let mut rows_checked = 0;
3138        for _ in 0..batches {
3139            let rows: [[f64; K]; 4] =
3140                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3141            let u: [[f64; K]; 4] =
3142                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3143            let v: [[f64; K]; 4] =
3144                std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3145
3146            let prod: [TwoSeed<K>; 4] = std::array::from_fn(|r| {
3147                let p: [TwoSeed<K>; K] =
3148                    std::array::from_fn(|a| TwoSeed::seed(rows[r][a], a, u[r][a], v[r][a]));
3149                row_expr(&p)
3150            });
3151
3152            let scal: [TwoSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
3153                let p: [TwoSeedLane<f64, K>; K] =
3154                    std::array::from_fn(|a| TwoSeedLane::seed(rows[r][a], a, u[r][a], v[r][a]));
3155                row_expr(&p)
3156            });
3157
3158            let pbatch: [TwoSeedBatch<K>; K] = std::array::from_fn(|a| {
3159                let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
3160                let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
3161                let vv = wide::f64x4::new([v[0][a], v[1][a], v[2][a], v[3][a]]);
3162                TwoSeedBatch::seed(val, a, uu, vv)
3163            });
3164            let batch = row_expr(&pbatch);
3165
3166            for r in 0..4 {
3167                let want = prod[r].contracted_fourth();
3168                let got_scal = scal[r].contracted_fourth();
3169                let got_batch = batch.lane(r).contracted_fourth();
3170                assert_eq!(
3171                    scal[r].base.v.to_bits(),
3172                    prod[r].base.value().to_bits(),
3173                    "TwoSeed K={K} scalar value"
3174                );
3175                assert_eq!(
3176                    batch.lane(r).base.value().to_bits(),
3177                    prod[r].base.value().to_bits(),
3178                    "TwoSeed K={K} batch lane {r} value"
3179                );
3180                for a in 0..K {
3181                    for b in 0..K {
3182                        assert_eq!(
3183                            got_scal[a][b].to_bits(),
3184                            want[a][b].to_bits(),
3185                            "TwoSeed K={K} scalar fourth[{a}][{b}]"
3186                        );
3187                        assert_eq!(
3188                            got_batch[a][b].to_bits(),
3189                            want[a][b].to_bits(),
3190                            "TwoSeed K={K} batch lane {r} fourth[{a}][{b}]"
3191                        );
3192                    }
3193                }
3194                rows_checked += 1;
3195            }
3196        }
3197        rows_checked
3198    }
3199
3200    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: the
3201    /// contracted-third channel of every `OneSeedLane` lane is `to_bits`-identical
3202    /// to the production [`OneSeed`] per row.
3203    #[test]
3204    fn oneseed_lanes_contracted_third_bit_identical() {
3205        let mut state = 0x1234_5678_9ABC_DEF0_u64;
3206        let batches = 2000;
3207        let rows_checked = check_oneseed::<2>(&mut state, batches)
3208            + check_oneseed::<3>(&mut state, batches)
3209            + check_oneseed::<4>(&mut state, batches)
3210            + check_oneseed::<9>(&mut state, batches);
3211        // 4 widths × `batches` batches × 4 rows each: a silently empty inner
3212        // loop would leave this at zero instead of passing as a no-op.
3213        assert_eq!(rows_checked, 4 * batches * 4);
3214    }
3215
3216    /// ≥2000 random 4-row batches per K, across K ∈ {2,3,4,9}: the
3217    /// contracted-fourth channel of every `TwoSeedLane` lane is `to_bits`-identical
3218    /// to the production [`TwoSeed`] per row.
3219    #[test]
3220    fn twoseed_lanes_contracted_fourth_bit_identical() {
3221        let mut state = 0x0FED_CBA9_8765_4321_u64;
3222        let batches = 2000;
3223        let rows_checked = check_twoseed::<2>(&mut state, batches)
3224            + check_twoseed::<3>(&mut state, batches)
3225            + check_twoseed::<4>(&mut state, batches)
3226            + check_twoseed::<9>(&mut state, batches);
3227        // 4 widths × `batches` batches × 4 rows each: a silently empty inner
3228        // loop would leave this at zero instead of passing as a no-op.
3229        assert_eq!(rows_checked, 4 * batches * 4);
3230    }
3231}
3232
3233#[cfg(test)]
3234mod unit_tests {
3235    use super::{JetScalar, Order1, Order2, filtered_implicit_solve_scalar};
3236
3237    // ── Order2 direct property tests ─────────────────────────────────────────
3238
3239    /// `Order2::constant(c)` carries value `c` and zero everywhere else.
3240    #[test]
3241    fn order2_constant_has_zero_derivatives() {
3242        let s = Order2::<3>::constant(7.5);
3243        assert_eq!(s.value(), 7.5);
3244        for a in 0..3 {
3245            assert_eq!(s.g()[a], 0.0, "grad[{a}] should be zero");
3246            for b in 0..3 {
3247                assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
3248            }
3249        }
3250    }
3251
3252    /// `Order2::variable(x, axis)` has unit gradient in slot `axis` and zero Hessian.
3253    #[test]
3254    fn order2_variable_has_unit_gradient_in_seeded_slot() {
3255        let x = -2.5_f64;
3256        let s = Order2::<4>::variable(x, 2);
3257        assert_eq!(s.value(), x);
3258        for a in 0..4 {
3259            let expected_g = if a == 2 { 1.0 } else { 0.0 };
3260            assert_eq!(s.g()[a], expected_g, "grad[{a}]");
3261            for b in 0..4 {
3262                assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
3263            }
3264        }
3265    }
3266
3267    /// `Order2::add` sums gradient channels; `sub` is the inverse on gradients.
3268    /// Uses integer-valued primaries so the value roundtrip is also exact.
3269    #[test]
3270    fn order2_add_sub_roundtrip() {
3271        let p = Order2::<2>::variable(3.0, 0);
3272        let q = Order2::<2>::variable(2.0, 1);
3273        let pq = JetScalar::add(&p, &q);
3274        // value = 3 + 2 = 5
3275        assert_eq!(pq.value(), 5.0, "add value");
3276        let back = JetScalar::sub(&pq, &q);
3277        // (p + q) - q gradient should equal p's gradient exactly
3278        for a in 0..2 {
3279            assert_eq!(back.g()[a], p.g()[a], "grad[{a}] roundtrip");
3280        }
3281    }
3282
3283    /// `Order2::mul` of two variables satisfies the Leibniz product rule:
3284    ///   ∂(p·q)/∂p = q,  ∂(p·q)/∂q = p,  ∂²(p·q)/∂p∂q = 1.
3285    #[test]
3286    fn order2_mul_satisfies_leibniz_rule() {
3287        let pv = 3.0_f64;
3288        let qv = -2.0_f64;
3289        let p = Order2::<2>::variable(pv, 0);
3290        let q = Order2::<2>::variable(qv, 1);
3291        let pq = JetScalar::mul(&p, &q);
3292        assert_eq!(pq.value(), pv * qv, "value = p·q");
3293        assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
3294        assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
3295        assert_eq!(pq.h()[0][1], 1.0, "∂²(p·q)/∂p∂q = 1");
3296        assert_eq!(pq.h()[1][0], 1.0, "∂²(p·q)/∂q∂p = 1 (symmetric)");
3297        assert_eq!(pq.h()[0][0], 0.0, "∂²(p·q)/∂p² = 0");
3298        assert_eq!(pq.h()[1][1], 0.0, "∂²(p·q)/∂q² = 0");
3299    }
3300
3301    /// `Order2::scale(s)` multiplies every channel by `s`.
3302    #[test]
3303    fn order2_scale_multiplies_all_channels() {
3304        let p = Order2::<2>::variable(4.0, 0);
3305        let s = 2.5_f64;
3306        let ps = JetScalar::scale(&p, s);
3307        assert_eq!(ps.value(), 4.0 * s);
3308        assert_eq!(ps.g()[0], 1.0 * s);
3309        assert_eq!(ps.g()[1], 0.0);
3310    }
3311
3312    /// `Order2::exp` at a constant has value `e^c`, gradient `e^c * g`, Hessian `e^c * (g⊗g + H)`.
3313    /// At a seeded variable `p₀`, the first derivative is `e^{p₀}` and second is `e^{p₀}`.
3314    #[test]
3315    fn order2_exp_derivative_stack_correct() {
3316        let p0 = 1.0_f64;
3317        let p = Order2::<1>::variable(p0, 0);
3318        let ep = JetScalar::exp(&p);
3319        let e = p0.exp();
3320        assert!((ep.value() - e).abs() < 1e-15, "exp value");
3321        assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p) = exp(p)");
3322        assert!((ep.h()[0][0] - e).abs() < 1e-15, "d²/dp² exp(p) = exp(p)");
3323    }
3324
3325    /// `Order2::ln` at a seeded variable: d/dp ln(p) = 1/p, d²/dp² ln(p) = -1/p².
3326    #[test]
3327    fn order2_ln_derivative_stack_correct() {
3328        let p0 = 2.0_f64;
3329        let p = Order2::<1>::variable(p0, 0);
3330        let lnp = JetScalar::ln(&p);
3331        assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
3332        assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
3333        assert!(
3334            (lnp.h()[0][0] - (-1.0 / (p0 * p0))).abs() < 1e-15,
3335            "d²/dp² ln(p) = -1/p²"
3336        );
3337    }
3338
3339    /// `exp` and `ln` are mutual inverses: `ln(exp(p)).value() == p` at the scalar.
3340    #[test]
3341    fn order2_exp_ln_roundtrip_at_value() {
3342        let p0 = 0.8_f64;
3343        let p = Order2::<1>::variable(p0, 0);
3344        let roundtrip = JetScalar::ln(&JetScalar::exp(&p));
3345        assert!((roundtrip.value() - p0).abs() < 1e-14, "ln(exp(p)) ≈ p");
3346    }
3347
3348    // ── Order1 tests ─────────────────────────────────────────────────────────
3349
3350    /// `Order1::constant` carries the correct value with all-zero gradient.
3351    #[test]
3352    fn order1_constant_has_zero_gradient() {
3353        let s = Order1::<3>::constant(-5.0);
3354        assert_eq!(s.value(), -5.0);
3355        for a in 0..3 {
3356            assert_eq!(s.g()[a], 0.0, "g[{a}] should be zero");
3357        }
3358    }
3359
3360    /// `Order1::variable(x, axis)` has unit gradient only in `axis`.
3361    #[test]
3362    fn order1_variable_has_unit_gradient_in_seeded_slot() {
3363        let s = Order1::<3>::variable(2.0, 1);
3364        assert_eq!(s.value(), 2.0);
3365        assert_eq!(s.g()[0], 0.0);
3366        assert_eq!(s.g()[1], 1.0);
3367        assert_eq!(s.g()[2], 0.0);
3368    }
3369
3370    /// `Order1::mul` satisfies the product rule (value and gradient, no Hessian).
3371    #[test]
3372    fn order1_mul_satisfies_product_rule() {
3373        let pv = 3.0_f64;
3374        let qv = -2.0_f64;
3375        let p = Order1::<2>::variable(pv, 0);
3376        let q = Order1::<2>::variable(qv, 1);
3377        let pq = JetScalar::mul(&p, &q);
3378        assert_eq!(pq.value(), pv * qv);
3379        assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
3380        assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
3381    }
3382
3383    /// `Order1::exp` carries the correct value and gradient `e^{p₀}`.
3384    #[test]
3385    fn order1_exp_has_correct_value_and_gradient() {
3386        let p0 = 0.5_f64;
3387        let p = Order1::<2>::variable(p0, 0);
3388        let ep = JetScalar::exp(&p);
3389        let e = p0.exp();
3390        assert!((ep.value() - e).abs() < 1e-15, "exp value");
3391        assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p)");
3392        assert_eq!(ep.g()[1], 0.0, "irrelevant gradient slot is zero");
3393    }
3394
3395    /// `Order1` and `Order2` agree on value and gradient for the same expression.
3396    #[test]
3397    fn order1_and_order2_agree_on_value_and_gradient() {
3398        let p0 = 1.3_f64;
3399        let q0 = -0.7_f64;
3400        // evaluate (p * q + p).exp() at (p0, q0)
3401        let p1 = Order1::<2>::variable(p0, 0);
3402        let q1 = Order1::<2>::variable(q0, 1);
3403        let expr1 = JetScalar::exp(&JetScalar::add(&JetScalar::mul(&p1, &q1), &p1));
3404
3405        let p2 = Order2::<2>::variable(p0, 0);
3406        let q2 = Order2::<2>::variable(q0, 1);
3407        let expr2 = JetScalar::exp(&JetScalar::add(&JetScalar::mul(&p2, &q2), &p2));
3408
3409        assert!(
3410            (expr1.value() - expr2.value()).abs() < 1e-14,
3411            "value mismatch"
3412        );
3413        for a in 0..2 {
3414            assert!(
3415                (expr1.g()[a] - expr2.g()[a]).abs() < 1e-14,
3416                "gradient[{a}] mismatch"
3417            );
3418        }
3419    }
3420
3421    // ── filtered_implicit_solve_scalar ────────────────────────────────────────
3422
3423    /// Lift the trivial linear constraint F(a, θ) = a - θ = 0 through `Order2<1>`.
3424    /// The exact lifted jet is a(θ) = θ, so value=θ₀, gradient=1.
3425    #[test]
3426    fn filtered_implicit_solve_linear_constraint_gives_exact_jet() {
3427        let theta0 = 3.0_f64;
3428        let theta = Order2::<1>::variable(theta0, 0);
3429        // a0 = theta0, F_a = 1, inv_fa = 1; 2 iters suffice for Order2.
3430        let a = filtered_implicit_solve_scalar::<1, Order2<1>>(theta0, 1.0, 2, |a_jet| {
3431            JetScalar::sub(a_jet, &theta)
3432        });
3433        assert!((a.value() - theta0).abs() < 1e-14, "value = theta0");
3434        // da/dtheta = 1 (identity)
3435        assert!((a.g()[0] - 1.0).abs() < 1e-14, "gradient = 1");
3436        // d²a/dtheta² = 0 (linear)
3437        assert!(a.h()[0][0].abs() < 1e-14, "hessian = 0");
3438    }
3439
3440    /// `filtered_implicit_solve_scalar` on a quadratic constraint F(a,θ)=a²-θ=0
3441    /// with primal root a₀=√θ₀, giving da/dθ = 1/(2√θ₀), d²a/dθ² = -1/(4θ₀^{3/2}).
3442    #[test]
3443    fn filtered_implicit_solve_quadratic_constraint_matches_analytic_derivatives() {
3444        let theta0 = 4.0_f64;
3445        let a0 = theta0.sqrt();
3446        let inv_fa = 1.0 / (2.0 * a0);
3447        let theta = Order2::<1>::variable(theta0, 0);
3448        // F(a,theta) = a*a - theta
3449        let a = filtered_implicit_solve_scalar::<1, Order2<1>>(a0, inv_fa, 2, |a_jet| {
3450            let aa = JetScalar::mul(a_jet, a_jet);
3451            JetScalar::sub(&aa, &theta)
3452        });
3453        let tol = 1e-12;
3454        assert!((a.value() - a0).abs() < tol, "value = sqrt(theta0)");
3455        let expected_g = 0.5 / a0;
3456        assert!(
3457            (a.g()[0] - expected_g).abs() < tol,
3458            "da/dtheta = 1/(2*sqrt)"
3459        );
3460        let expected_h = -0.25 / (theta0 * a0);
3461        assert!(
3462            (a.h()[0][0] - expected_h).abs() < tol,
3463            "d2a/dtheta2 = -1/(4*theta^1.5)"
3464        );
3465    }
3466}