1pub trait JetScalar<const K: usize>: Copy {
73 fn constant(c: f64) -> Self;
75
76 fn variable(x: f64, axis: usize) -> Self;
81
82 fn value(&self) -> f64;
84
85 fn add(&self, o: &Self) -> Self;
87 fn sub(&self, o: &Self) -> Self;
89 fn mul(&self, o: &Self) -> Self;
91 fn neg(&self) -> Self;
93 fn scale(&self, s: f64) -> Self;
95
96 fn compose_unary(&self, d: [f64; 5]) -> Self;
107
108 fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
125 self.compose_unary(stack_fn(self.value()))
126 }
127
128 fn exp(&self) -> Self {
130 let e = self.value().exp();
131 self.compose_unary([e, e, e, e, e])
132 }
133
134 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 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 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 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 fn ln_gamma(&self) -> Self {
181 self.compose_unary(crate::jet_tower::ln_gamma_derivative_stack(self.value()))
182 }
183
184 fn digamma(&self) -> Self {
188 self.compose_unary(crate::jet_tower::digamma_derivative_stack(self.value()))
189 }
190}
191
192pub trait RuntimeJetScalar<'arena>: Clone {
200 type Workspace: ?Sized;
203
204 fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
206 fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
208 fn dimension(&self) -> usize;
210 fn value(&self) -> f64;
212 fn add(&self, o: &Self) -> Self;
214 fn sub(&self, o: &Self) -> Self;
216 fn mul(&self, o: &Self) -> Self;
218 fn neg(&self) -> Self;
220 fn scale(&self, s: f64) -> Self;
222 fn compose_unary(&self, d: [f64; 5]) -> Self;
224
225 fn exp(&self) -> Self {
227 let e = self.value().exp();
228 self.compose_unary([e, e, e, e, e])
229 }
230
231 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#[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 #[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#[derive(Debug)]
322pub struct DynamicJetArena {
323 bump: bumpalo::Bump,
324}
325
326impl DynamicJetArena {
327 #[must_use]
329 pub fn new() -> Self {
330 Self {
331 bump: bumpalo::Bump::new(),
332 }
333 }
334
335 #[must_use]
337 pub fn with_capacity(bytes: usize) -> Self {
338 Self {
339 bump: bumpalo::Bump::with_capacity(bytes),
340 }
341 }
342
343 pub fn reset(&mut self) {
345 self.bump.reset();
346 }
347
348 #[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 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#[derive(Clone, Copy, Debug)]
376pub struct DynamicOrder1<'arena> {
377 arena: &'arena DynamicJetArena,
378 pub v: f64,
380 pub g: &'arena [f64],
382}
383
384impl DynamicOrder1<'_> {
385 #[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#[derive(Clone, Copy, Debug)]
506pub struct DynamicOrder2<'arena> {
507 arena: &'arena DynamicJetArena,
508 pub v: f64,
510 pub g: &'arena [f64],
512 pub h: &'arena [f64],
514}
515
516impl DynamicOrder2<'_> {
517 #[inline]
519 #[must_use]
520 pub fn g(&self) -> &[f64] {
521 self.g
522 }
523
524 #[inline]
526 #[must_use]
527 pub fn h(&self) -> &[f64] {
528 self.h
529 }
530
531 #[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#[derive(Clone, Copy, Debug)]
698pub struct DynamicOneSeed<'arena> {
699 pub base: DynamicOrder2<'arena>,
701 pub eps: DynamicOrder2<'arena>,
703}
704
705impl<'arena> DynamicOneSeed<'arena> {
706 #[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 #[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#[derive(Clone, Copy, Debug)]
798pub struct DynamicTwoSeed<'arena> {
799 pub base: DynamicOrder2<'arena>,
801 pub eps: DynamicOrder2<'arena>,
803 pub del: DynamicOrder2<'arena>,
805 pub eps_del: DynamicOrder2<'arena>,
807}
808
809impl<'arena> DynamicTwoSeed<'arena> {
810 #[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 #[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
937impl<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
1002pub 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#[derive(Clone, Copy, Debug)]
1053pub struct Order2<const K: usize>(pub crate::jet_tower::Tower2<K>);
1054
1055impl<const K: usize> Order2<K> {
1056 #[inline]
1058 pub fn g(&self) -> [f64; K] {
1059 self.0.g
1060 }
1061
1062 #[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 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 Order2(self.0.compose_unary([d[0], d[1], d[2]]))
1099 }
1100}
1101
1102pub trait Lane: Copy {
1138 fn splat(x: f64) -> Self;
1140 fn add(self, o: Self) -> Self;
1142 fn sub(self, o: Self) -> Self;
1144 fn mul(self, o: Self) -> Self;
1146 fn lane(self, i: usize) -> f64;
1148 fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3];
1154 fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5];
1163 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 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 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#[derive(Clone, Copy, Debug)]
1297pub struct Order2Lane<L: Lane, const K: usize> {
1298 pub v: L,
1300 pub g: [L; K],
1302 pub h: [[L; K]; K],
1304}
1305
1306pub type Order2Batch<const K: usize> = Order2Lane<wide::f64x4, K>;
1308
1309impl<L: Lane, const K: usize> Order2Lane<L, K> {
1310 #[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 #[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 #[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 #[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 #[inline]
1362 pub fn sub(&self, o: &Self) -> Self {
1363 self.add(&o.scale(-1.0))
1364 }
1365
1366 #[inline]
1368 pub fn neg(&self) -> Self {
1369 self.scale(-1.0)
1370 }
1371
1372 #[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 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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#[derive(Clone, Copy, Debug)]
1526pub struct Order1<const K: usize> {
1527 pub v: f64,
1529 pub g: [f64; K],
1531}
1532
1533impl<const K: usize> Order1<K> {
1534 #[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 Order1 { v: c, g: [0.0; K] }
1545 }
1546 fn variable(x: f64, axis: usize) -> Self {
1547 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 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 self.add(&o.scale(-1.0))
1566 }
1567 fn mul(&self, o: &Self) -> Self {
1568 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 self.scale(-1.0)
1583 }
1584 fn scale(&self, s: f64) -> Self {
1585 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 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#[derive(Clone, Copy, Debug)]
1623pub struct OneSeed<const K: usize> {
1624 pub base: Order2<K>,
1626 pub eps: Order2<K>,
1629}
1630
1631impl<const K: usize> OneSeed<K> {
1632 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 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 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 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 let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
1705 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#[derive(Clone, Copy, Debug)]
1724pub struct OneSeedLane<L: Lane, const K: usize> {
1725 pub base: Order2Lane<L, K>,
1727 pub eps: Order2Lane<L, K>,
1730}
1731
1732pub type OneSeedBatch<const K: usize> = OneSeedLane<wide::f64x4, K>;
1734
1735impl<L: Lane, const K: usize> OneSeedLane<L, K> {
1736 #[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 #[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 #[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 #[inline]
1770 #[must_use]
1771 pub fn contracted_third(&self) -> [[L; K]; K] {
1772 self.eps.h
1773 }
1774
1775 #[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 #[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 #[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 #[inline]
1804 pub fn neg(&self) -> Self {
1805 OneSeedLane {
1806 base: self.base.neg(),
1807 eps: self.eps.neg(),
1808 }
1809 }
1810
1811 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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#[derive(Clone, Copy, Debug)]
1949pub struct TwoSeed<const K: usize> {
1950 pub base: Order2<K>,
1952 pub eps: Order2<K>,
1954 pub del: Order2<K>,
1956 pub eps_del: Order2<K>,
1959}
1960
1961impl<const K: usize> TwoSeed<K> {
1962 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 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 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 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]]); let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]); 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#[derive(Clone, Copy, Debug)]
2092pub struct TwoSeedLane<L: Lane, const K: usize> {
2093 pub base: Order2Lane<L, K>,
2095 pub eps: Order2Lane<L, K>,
2097 pub del: Order2Lane<L, K>,
2099 pub eps_del: Order2Lane<L, K>,
2102}
2103
2104pub type TwoSeedBatch<const K: usize> = TwoSeedLane<wide::f64x4, K>;
2106
2107impl<L: Lane, const K: usize> TwoSeedLane<L, K> {
2108 #[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 #[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 #[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 #[inline]
2151 #[must_use]
2152 pub fn contracted_fourth(&self) -> [[L; K]; K] {
2153 self.eps_del.h
2154 }
2155
2156 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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
2347impl<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
2383impl<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 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 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 #[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 #[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 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 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 #[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 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 #[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 #[test]
2618 fn generic_program_seam_matches_tower_for_every_channel() {
2619 let t = tower();
2620 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 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 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 #[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 #[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 use super::{
2775 JetScalar, Lane, OneSeed, OneSeedBatch, OneSeedLane, Order2, Order2Batch, Order2Lane,
2776 TwoSeed, TwoSeedBatch, TwoSeedLane,
2777 };
2778
2779 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 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 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); 2.0 * u - 1.0
2872 }
2873
2874 fn check_k<const K: usize>(state: &mut u64, batches: usize) -> usize {
2877 let mut verified_rows = 0usize;
2878 for _ in 0..batches {
2879 let rows: [[f64; K]; 4] =
2881 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
2882
2883 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 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 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 assert_eq!(scal[r].v.to_bits(), g.v.to_bits(), "K={K} scalar v");
2907 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 #[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 assert_eq!(verified, 4 * 2000 * 4, "every batch row must be verified");
2952 }
2953
2954 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 let u: [[f64; K]; 4] =
3077 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
3078
3079 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 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 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 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 #[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 assert_eq!(rows_checked, 4 * batches * 4);
3214 }
3215
3216 #[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 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 #[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 #[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 #[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 assert_eq!(pq.value(), 5.0, "add value");
3276 let back = JetScalar::sub(&pq, &q);
3277 for a in 0..2 {
3279 assert_eq!(back.g()[a], p.g()[a], "grad[{a}] roundtrip");
3280 }
3281 }
3282
3283 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[test]
3397 fn order1_and_order2_agree_on_value_and_gradient() {
3398 let p0 = 1.3_f64;
3399 let q0 = -0.7_f64;
3400 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 #[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 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 assert!((a.g()[0] - 1.0).abs() < 1e-14, "gradient = 1");
3436 assert!(a.h()[0][0].abs() < 1e-14, "hessian = 0");
3438 }
3439
3440 #[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 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}