1pub trait SymmetricQuadraticCoefficients {
72 fn dimension(&self) -> usize;
74
75 fn multiply(&self, input: &[f64], output: &mut [f64]);
77
78 fn coefficient(&self, row: usize, column: usize) -> f64;
80
81 fn visit_upper_triangle(
92 &self,
93 direction: &mut [f64],
94 projected: &mut [f64],
95 mut visit: impl FnMut(usize, usize, f64),
96 ) {
97 let dimension = self.dimension();
98 assert_eq!(direction.len(), dimension);
99 assert_eq!(projected.len(), dimension);
100 direction.fill(0.0);
101 for column in 0..dimension {
102 direction[column] = 1.0;
103 self.multiply(direction, projected);
104 direction[column] = 0.0;
105 for row in 0..=column {
106 visit(row, column, projected[row]);
107 }
108 }
109 }
110
111 fn quadratic_value<T, F>(&self, inputs: &[T], value: F) -> f64
116 where
117 F: Fn(&T) -> f64,
118 {
119 assert_eq!(
120 inputs.len(),
121 self.dimension(),
122 "symmetric quadratic-form dimension mismatch"
123 );
124 let mut out = 0.0;
125 for row in 0..inputs.len() {
126 let row_value = value(&inputs[row]);
127 out += self.coefficient(row, row) * row_value * row_value;
128 for column in row + 1..inputs.len() {
129 out += 2.0 * self.coefficient(row, column) * row_value * value(&inputs[column]);
130 }
131 }
132 out
133 }
134}
135
136fn symmetric_quadratic_form_default<T, C>(
137 inputs: &[T],
138 coefficients: &C,
139 constant: impl Fn(f64) -> T,
140 add: impl Fn(&T, &T) -> T,
141 mul: impl Fn(&T, &T) -> T,
142 scale: impl Fn(&T, f64) -> T,
143) -> T
144where
145 C: SymmetricQuadraticCoefficients,
146{
147 assert_eq!(
148 inputs.len(),
149 coefficients.dimension(),
150 "symmetric quadratic-form dimension mismatch"
151 );
152 let mut out = constant(0.0);
153 for row in 0..inputs.len() {
154 let diagonal = mul(&inputs[row], &inputs[row]);
155 out = add(&out, &scale(&diagonal, coefficients.coefficient(row, row)));
156 for column in row + 1..inputs.len() {
157 let cross = mul(&inputs[row], &inputs[column]);
158 out = add(
159 &out,
160 &scale(&cross, 2.0 * coefficients.coefficient(row, column)),
161 );
162 }
163 }
164 out
165}
166
167fn linear_combination_default<T>(
168 inputs: &[T],
169 weights: &[f64],
170 constant: impl Fn(f64) -> T,
171 add: impl Fn(&T, &T) -> T,
172 scale: impl Fn(&T, f64) -> T,
173) -> T {
174 assert_eq!(
175 inputs.len(),
176 weights.len(),
177 "linear-combination dimension mismatch"
178 );
179 inputs
180 .iter()
181 .zip(weights)
182 .fold(constant(0.0), |sum, (input, &weight)| {
183 add(&sum, &scale(input, weight))
184 })
185}
186
187fn multiply_add_default<T>(
188 left: &T,
189 right: &T,
190 addend: &T,
191 mul: impl Fn(&T, &T) -> T,
192 add: impl Fn(&T, &T) -> T,
193) -> T {
194 add(&mul(left, right), addend)
195}
196
197fn composed_sum_default<T>(
198 inputs: &[T],
199 derivative_stacks: &[[f64; 5]],
200 constant: impl Fn(f64) -> T,
201 add: impl Fn(&T, &T) -> T,
202 compose: impl Fn(&T, [f64; 5]) -> T,
203) -> T {
204 assert_eq!(
205 inputs.len(),
206 derivative_stacks.len(),
207 "composed-sum term-count mismatch"
208 );
209 inputs
210 .iter()
211 .zip(derivative_stacks)
212 .fold(constant(0.0), |sum, (input, &stack)| {
213 add(&sum, &compose(input, stack))
214 })
215}
216
217fn affine_compose_default<T>(
218 input: &T,
219 input_scale: f64,
220 input_shift: f64,
221 derivative_stack: [f64; 5],
222 scale: impl Fn(&T, f64) -> T,
223 add_constant: impl Fn(&T, f64) -> T,
224 compose: impl Fn(&T, [f64; 5]) -> T,
225) -> T {
226 compose(
227 &add_constant(&scale(input, input_scale), input_shift),
228 derivative_stack,
229 )
230}
231
232fn affine_composed_sum_default<T>(
233 inputs: &[T],
234 input_scales: &[f64],
235 derivative_stacks: &[[f64; 5]],
236 constant: impl Fn(f64) -> T,
237 add: impl Fn(&T, &T) -> T,
238 scale: impl Fn(&T, f64) -> T,
239 add_constant: impl Fn(&T, f64) -> T,
240 compose: impl Fn(&T, [f64; 5]) -> T,
241) -> T {
242 assert_eq!(inputs.len(), input_scales.len());
243 assert_eq!(inputs.len(), derivative_stacks.len());
244 inputs.iter().zip(input_scales).zip(derivative_stacks).fold(
245 constant(0.0),
246 |sum, ((input, &input_scale), &stack)| {
247 add(
248 &sum,
249 &affine_compose_default(
250 input,
251 input_scale,
252 0.0,
253 stack,
254 &scale,
255 &add_constant,
256 &compose,
257 ),
258 )
259 },
260 )
261}
262
263fn shared_multiply_add_affine_composed_sum_default<T, const N: usize>(
264 lefts: &[&T; N],
265 right: &T,
266 addend: &T,
267 addend_scales: &[f64; N],
268 input_scales: &[f64; N],
269 derivative_stacks: &[[f64; 5]; N],
270 constant: impl Fn(f64) -> T,
271 add: impl Fn(&T, &T) -> T,
272 mul: impl Fn(&T, &T) -> T,
273 scale: impl Fn(&T, f64) -> T,
274 multiply_add: impl Fn(&T, &T, &T) -> T,
275 affine_compose: impl Fn(&T, f64, f64, [f64; 5]) -> T,
276) -> T {
277 let (representatives, term_sources, source_count) =
278 canonical_shared_source_schedule(|term, representative| {
279 std::ptr::eq(lefts[term], lefts[representative])
280 && addend_scales[term] == addend_scales[representative]
281 });
282 let (value, source_derivatives) =
283 aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
284 (0..source_count).fold(constant(value), |sum, source| {
285 let term = representatives[source];
286 let inner = if addend_scales[term] == 0.0 {
287 mul(lefts[term], right)
288 } else if addend_scales[term] == 1.0 {
289 multiply_add(lefts[term], right, addend)
290 } else {
291 multiply_add(lefts[term], right, &scale(addend, addend_scales[term]))
292 };
293 let composed = affine_compose(&inner, 1.0, 0.0, source_derivatives[source]);
294 add(&sum, &composed)
295 })
296}
297
298#[inline(always)]
306pub(crate) fn canonical_shared_source_schedule<const N: usize>(
307 mut equivalent: impl FnMut(usize, usize) -> bool,
308) -> ([usize; N], [usize; N], usize) {
309 let mut representatives = [0; N];
310 let mut term_sources = [0; N];
311 let mut source_count = 0;
312 for term in 0..N {
313 let mut source = 0;
314 while source < source_count && !equivalent(term, representatives[source]) {
315 source += 1;
316 }
317 if source == source_count {
318 representatives[source] = term;
319 source_count += 1;
320 }
321 term_sources[term] = source;
322 }
323 (representatives, term_sources, source_count)
324}
325
326#[inline(always)]
333pub(crate) fn aggregate_shared_source_derivatives<const N: usize>(
334 term_sources: &[usize; N],
335 input_scales: &[f64; N],
336 derivative_stacks: &[[f64; 5]; N],
337) -> (f64, [[f64; 5]; N]) {
338 let mut value = 0.0;
339 let mut source_derivatives = [[0.0; 5]; N];
340 for term in 0..N {
341 value += derivative_stacks[term][0];
342 let source = term_sources[term];
343 let mut scale_power = input_scales[term];
344 for order in 1..5 {
345 source_derivatives[source][order] += derivative_stacks[term][order] * scale_power;
346 scale_power *= input_scales[term];
347 }
348 }
349 (value, source_derivatives)
350}
351
352pub trait JetScalar<const K: usize>: crate::nested_dual::JetField + Copy {
360 fn constant(c: f64) -> Self;
362
363 fn variable(x: f64, axis: usize) -> Self;
368
369 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
373 inputs: &[Self],
374 coefficients: &C,
375 ) -> Self {
376 symmetric_quadratic_form_default(
377 inputs,
378 coefficients,
379 Self::constant,
380 crate::nested_dual::JetField::add,
381 crate::nested_dual::JetField::mul,
382 crate::nested_dual::JetField::scale,
383 )
384 }
385
386 fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
388 linear_combination_default(
389 inputs,
390 weights,
391 Self::constant,
392 crate::nested_dual::JetField::add,
393 crate::nested_dual::JetField::scale,
394 )
395 }
396
397 fn add_constant(&self, constant: f64) -> Self {
399 self.add(&Self::constant(constant))
400 }
401
402 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
404 multiply_add_default(
405 self,
406 right,
407 addend,
408 crate::nested_dual::JetField::mul,
409 crate::nested_dual::JetField::add,
410 )
411 }
412
413 fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
415 composed_sum_default(
416 inputs,
417 derivative_stacks,
418 Self::constant,
419 crate::nested_dual::JetField::add,
420 crate::nested_dual::JetField::compose_unary,
421 )
422 }
423
424 fn product(&self, right: &Self) -> Self {
426 self.mul(right)
427 }
428
429 fn affine_compose(
432 &self,
433 input_scale: f64,
434 input_shift: f64,
435 derivative_stack: [f64; 5],
436 ) -> Self {
437 affine_compose_default(
438 self,
439 input_scale,
440 input_shift,
441 derivative_stack,
442 crate::nested_dual::JetField::scale,
443 Self::add_constant,
444 crate::nested_dual::JetField::compose_unary,
445 )
446 }
447
448 fn affine_composed_sum(
450 inputs: &[Self],
451 input_scales: &[f64],
452 derivative_stacks: &[[f64; 5]],
453 ) -> Self {
454 affine_composed_sum_default(
455 inputs,
456 input_scales,
457 derivative_stacks,
458 Self::constant,
459 crate::nested_dual::JetField::add,
460 crate::nested_dual::JetField::scale,
461 Self::add_constant,
462 crate::nested_dual::JetField::compose_unary,
463 )
464 }
465
466 fn shared_multiply_add_affine_composed_sum<const N: usize>(
475 lefts: &[&Self; N],
476 right: &Self,
477 addend: &Self,
478 addend_scales: &[f64; N],
479 input_scales: &[f64; N],
480 derivative_stacks: &[[f64; 5]; N],
481 ) -> Self {
482 shared_multiply_add_affine_composed_sum_default(
483 lefts,
484 right,
485 addend,
486 addend_scales,
487 input_scales,
488 derivative_stacks,
489 Self::constant,
490 crate::nested_dual::JetField::add,
491 crate::nested_dual::JetField::mul,
492 crate::nested_dual::JetField::scale,
493 Self::multiply_add,
494 Self::affine_compose,
495 )
496 }
497
498 fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
511 self.compose_unary(stack_fn(self.value()))
512 }
513
514 fn exp(&self) -> Self {
516 let e = self.value().exp();
517 self.compose_unary([e, e, e, e, e])
518 }
519
520 fn sqrt(&self) -> Self {
522 let u = self.value();
523 let s = u.sqrt();
524 self.compose_unary([
525 s,
526 0.5 / s,
527 -0.25 / (u * s),
528 0.375 / (u * u * s),
529 -0.9375 / (u * u * u * s),
530 ])
531 }
532
533 fn ln(&self) -> Self {
537 let u = self.value();
538 let r = 1.0 / u;
539 self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
540 }
541
542 fn recip(&self) -> Self {
544 let r = 1.0 / self.value();
545 let r2 = r * r;
546 self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
547 }
548
549 fn powf(&self, a: f64) -> Self {
552 let u = self.value();
553 self.compose_unary([
554 u.powf(a),
555 a * u.powf(a - 1.0),
556 a * (a - 1.0) * u.powf(a - 2.0),
557 a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
558 a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
559 ])
560 }
561
562 fn ln_gamma(&self) -> Self {
567 self.compose_unary(crate::jet_tower::ln_gamma_derivative_stack(self.value()))
568 }
569
570 fn digamma(&self) -> Self {
574 self.compose_unary(crate::jet_tower::digamma_derivative_stack(self.value()))
575 }
576}
577
578impl<S, const K: usize> JetScalar<K> for crate::nested_dual::Dual2<S>
583where
584 S: JetScalar<K>,
585{
586 #[inline]
587 fn constant(c: f64) -> Self {
588 Self {
589 v: S::constant(c),
590 g: S::constant(0.0),
591 h: S::constant(0.0),
592 }
593 }
594
595 #[inline]
596 fn variable(x: f64, axis: usize) -> Self {
597 Self {
598 v: S::variable(x, axis),
599 g: S::constant(0.0),
600 h: S::constant(0.0),
601 }
602 }
603}
604
605pub trait RuntimeJetScalar<'arena>: Clone {
613 type Workspace: ?Sized;
616
617 fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
619 fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self;
621
622 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
625 inputs: &[Self],
626 coefficients: &C,
627 dimension: usize,
628 workspace: &'arena Self::Workspace,
629 ) -> Self {
630 symmetric_quadratic_form_default(
631 inputs,
632 coefficients,
633 |value| Self::constant(value, dimension, workspace),
634 Self::add,
635 Self::mul,
636 Self::scale,
637 )
638 }
639
640 fn linear_combination(
642 inputs: &[Self],
643 weights: &[f64],
644 dimension: usize,
645 workspace: &'arena Self::Workspace,
646 ) -> Self {
647 linear_combination_default(
648 inputs,
649 weights,
650 |value| Self::constant(value, dimension, workspace),
651 Self::add,
652 Self::scale,
653 )
654 }
655
656 fn add_constant(&self, constant: f64, workspace: &'arena Self::Workspace) -> Self {
658 self.add(&Self::constant(constant, self.dimension(), workspace))
659 }
660
661 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
663 multiply_add_default(self, right, addend, Self::mul, Self::add)
664 }
665
666 fn composed_sum(
668 inputs: &[Self],
669 derivative_stacks: &[[f64; 5]],
670 dimension: usize,
671 workspace: &'arena Self::Workspace,
672 ) -> Self {
673 composed_sum_default(
674 inputs,
675 derivative_stacks,
676 |value| Self::constant(value, dimension, workspace),
677 Self::add,
678 Self::compose_unary,
679 )
680 }
681
682 fn product(&self, right: &Self) -> Self {
684 self.mul(right)
685 }
686
687 fn affine_compose(
689 &self,
690 input_scale: f64,
691 input_shift: f64,
692 derivative_stack: [f64; 5],
693 workspace: &'arena Self::Workspace,
694 ) -> Self {
695 affine_compose_default(
696 self,
697 input_scale,
698 input_shift,
699 derivative_stack,
700 Self::scale,
701 |value, constant| value.add_constant(constant, workspace),
702 Self::compose_unary,
703 )
704 }
705
706 fn affine_composed_sum(
708 inputs: &[Self],
709 input_scales: &[f64],
710 derivative_stacks: &[[f64; 5]],
711 dimension: usize,
712 workspace: &'arena Self::Workspace,
713 ) -> Self {
714 affine_composed_sum_default(
715 inputs,
716 input_scales,
717 derivative_stacks,
718 |value| Self::constant(value, dimension, workspace),
719 Self::add,
720 Self::scale,
721 |value, constant| value.add_constant(constant, workspace),
722 Self::compose_unary,
723 )
724 }
725
726 fn shared_multiply_add_affine_composed_sum<const N: usize>(
734 lefts: &[&Self; N],
735 right: &Self,
736 addend: &Self,
737 addend_scales: &[f64; N],
738 input_scales: &[f64; N],
739 derivative_stacks: &[[f64; 5]; N],
740 dimension: usize,
741 workspace: &'arena Self::Workspace,
742 ) -> Self {
743 shared_multiply_add_affine_composed_sum_default(
744 lefts,
745 right,
746 addend,
747 addend_scales,
748 input_scales,
749 derivative_stacks,
750 |value| Self::constant(value, dimension, workspace),
751 Self::add,
752 Self::mul,
753 Self::scale,
754 Self::multiply_add,
755 |input, scale, shift, stack| input.affine_compose(scale, shift, stack, workspace),
756 )
757 }
758 fn dimension(&self) -> usize;
760 fn value(&self) -> f64;
762 fn add(&self, o: &Self) -> Self;
764 fn sub(&self, o: &Self) -> Self;
766 fn mul(&self, o: &Self) -> Self;
768 fn neg(&self) -> Self;
770 fn scale(&self, s: f64) -> Self;
772 fn compose_unary(&self, d: [f64; 5]) -> Self;
774
775 fn exp(&self) -> Self {
777 let e = self.value().exp();
778 self.compose_unary([e, e, e, e, e])
779 }
780
781 fn ln(&self) -> Self {
785 let u = self.value();
786 let r = 1.0 / u;
787 self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
788 }
789
790 fn recip(&self) -> Self {
792 let r = 1.0 / self.value();
793 let r2 = r * r;
794 self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
795 }
796}
797
798#[derive(Clone, Copy, Debug, PartialEq)]
806pub struct RuntimeValue {
807 value: f64,
808 dimension: usize,
809}
810
811impl<'arena> RuntimeJetScalar<'arena> for RuntimeValue {
812 type Workspace = ();
813
814 #[inline(always)]
815 fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
816 Self {
817 value: c,
818 dimension,
819 }
820 }
821
822 #[inline(always)]
823 fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
824 assert!(
825 axis < dimension,
826 "runtime value variable axis out of bounds"
827 );
828 Self {
829 value: x,
830 dimension,
831 }
832 }
833
834 #[inline(always)]
835 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
836 inputs: &[Self],
837 coefficients: &C,
838 dimension: usize,
839 &(): &'arena Self::Workspace,
840 ) -> Self {
841 assert_eq!(inputs.len(), coefficients.dimension());
842 assert!(inputs.iter().all(|input| input.dimension == dimension));
843 Self {
844 value: coefficients.quadratic_value(inputs, |input| input.value),
845 dimension,
846 }
847 }
848
849 #[inline(always)]
850 fn linear_combination(
851 inputs: &[Self],
852 weights: &[f64],
853 dimension: usize,
854 &(): &'arena Self::Workspace,
855 ) -> Self {
856 assert_eq!(inputs.len(), weights.len());
857 assert!(inputs.iter().all(|input| input.dimension == dimension));
858 let value = inputs
859 .iter()
860 .zip(weights)
861 .map(|(input, &weight)| input.value * weight)
862 .sum();
863 Self { value, dimension }
864 }
865
866 #[inline(always)]
867 fn add_constant(&self, constant: f64, &(): &'arena Self::Workspace) -> Self {
868 Self {
869 value: self.value + constant,
870 dimension: self.dimension,
871 }
872 }
873
874 #[inline(always)]
875 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
876 self.assert_same_dimension(right);
877 self.assert_same_dimension(addend);
878 Self {
879 value: self.value * right.value + addend.value,
880 dimension: self.dimension,
881 }
882 }
883
884 #[inline(always)]
885 fn composed_sum(
886 inputs: &[Self],
887 derivative_stacks: &[[f64; 5]],
888 dimension: usize,
889 &(): &'arena Self::Workspace,
890 ) -> Self {
891 assert_eq!(inputs.len(), derivative_stacks.len());
892 assert!(inputs.iter().all(|input| input.dimension == dimension));
893 Self {
894 value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
895 dimension,
896 }
897 }
898
899 #[inline(always)]
900 fn product(&self, right: &Self) -> Self {
901 self.mul(right)
902 }
903
904 #[inline(always)]
905 fn affine_compose(
906 &self,
907 _: f64,
908 _: f64,
909 derivative_stack: [f64; 5],
910 &(): &'arena Self::Workspace,
911 ) -> Self {
912 Self {
913 value: derivative_stack[0],
914 dimension: self.dimension,
915 }
916 }
917
918 #[inline(always)]
919 fn affine_composed_sum(
920 inputs: &[Self],
921 input_scales: &[f64],
922 derivative_stacks: &[[f64; 5]],
923 dimension: usize,
924 &(): &'arena Self::Workspace,
925 ) -> Self {
926 assert_eq!(inputs.len(), input_scales.len());
927 assert_eq!(inputs.len(), derivative_stacks.len());
928 assert!(inputs.iter().all(|input| input.dimension == dimension));
929 Self {
930 value: derivative_stacks.iter().map(|stack| stack[0]).sum(),
931 dimension,
932 }
933 }
934
935 #[inline(always)]
936 fn dimension(&self) -> usize {
937 self.dimension
938 }
939
940 #[inline(always)]
941 fn value(&self) -> f64 {
942 self.value
943 }
944
945 #[inline(always)]
946 fn add(&self, other: &Self) -> Self {
947 self.assert_same_dimension(other);
948 Self {
949 value: self.value + other.value,
950 dimension: self.dimension,
951 }
952 }
953
954 #[inline(always)]
955 fn sub(&self, other: &Self) -> Self {
956 self.assert_same_dimension(other);
957 Self {
958 value: self.value - other.value,
959 dimension: self.dimension,
960 }
961 }
962
963 #[inline(always)]
964 fn mul(&self, other: &Self) -> Self {
965 self.assert_same_dimension(other);
966 Self {
967 value: self.value * other.value,
968 dimension: self.dimension,
969 }
970 }
971
972 #[inline(always)]
973 fn neg(&self) -> Self {
974 Self {
975 value: -self.value,
976 dimension: self.dimension,
977 }
978 }
979
980 #[inline(always)]
981 fn scale(&self, scale: f64) -> Self {
982 Self {
983 value: self.value * scale,
984 dimension: self.dimension,
985 }
986 }
987
988 #[inline(always)]
989 fn compose_unary(&self, derivative_stack: [f64; 5]) -> Self {
990 Self {
991 value: derivative_stack[0],
992 dimension: self.dimension,
993 }
994 }
995}
996
997impl RuntimeValue {
998 #[inline(always)]
999 fn assert_same_dimension(&self, other: &Self) {
1000 assert_eq!(self.dimension, other.dimension);
1001 }
1002}
1003
1004#[derive(Clone, Copy, Debug)]
1009#[repr(transparent)]
1010pub struct FixedRuntimeJet<S, const K: usize> {
1011 inner: S,
1012}
1013
1014impl<S, const K: usize> FixedRuntimeJet<S, K> {
1015 #[inline(always)]
1018 #[must_use]
1019 pub fn from_inner(inner: S) -> Self {
1020 Self { inner }
1021 }
1022
1023 #[inline(always)]
1025 #[must_use]
1026 pub fn into_inner(self) -> S {
1027 self.inner
1028 }
1029}
1030
1031impl<'arena, S: JetScalar<K>, const K: usize> RuntimeJetScalar<'arena> for FixedRuntimeJet<S, K> {
1032 type Workspace = ();
1033
1034 #[inline(always)]
1035 fn constant(c: f64, dimension: usize, &(): &'arena Self::Workspace) -> Self {
1036 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1037 Self {
1038 inner: S::constant(c),
1039 }
1040 }
1041
1042 #[inline(always)]
1043 fn variable(x: f64, axis: usize, dimension: usize, &(): &'arena Self::Workspace) -> Self {
1044 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1045 Self {
1046 inner: S::variable(x, axis),
1047 }
1048 }
1049
1050 #[inline(always)]
1051 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
1052 inputs: &[Self],
1053 coefficients: &C,
1054 dimension: usize,
1055 &(): &'arena Self::Workspace,
1056 ) -> Self {
1057 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1058 assert_eq!(inputs.len(), coefficients.dimension());
1059 let inner =
1065 unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1066 Self {
1067 inner: S::symmetric_quadratic_form(inner, coefficients),
1068 }
1069 }
1070
1071 #[inline(always)]
1072 fn linear_combination(
1073 inputs: &[Self],
1074 weights: &[f64],
1075 dimension: usize,
1076 &(): &'arena Self::Workspace,
1077 ) -> Self {
1078 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1079 assert_eq!(inputs.len(), weights.len());
1080 let inner =
1084 unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1085 Self {
1086 inner: S::linear_combination(inner, weights),
1087 }
1088 }
1089
1090 #[inline(always)]
1091 fn add_constant(&self, constant: f64, &(): &'arena Self::Workspace) -> Self {
1092 Self {
1093 inner: self.inner.add_constant(constant),
1094 }
1095 }
1096
1097 #[inline(always)]
1098 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
1099 Self {
1100 inner: self.inner.multiply_add(&right.inner, &addend.inner),
1101 }
1102 }
1103
1104 #[inline(always)]
1105 fn composed_sum(
1106 inputs: &[Self],
1107 derivative_stacks: &[[f64; 5]],
1108 dimension: usize,
1109 &(): &'arena Self::Workspace,
1110 ) -> Self {
1111 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1112 let inner =
1116 unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1117 Self {
1118 inner: S::composed_sum(inner, derivative_stacks),
1119 }
1120 }
1121
1122 #[inline(always)]
1123 fn product(&self, right: &Self) -> Self {
1124 Self {
1125 inner: self.inner.product(&right.inner),
1126 }
1127 }
1128
1129 #[inline(always)]
1130 fn affine_compose(
1131 &self,
1132 input_scale: f64,
1133 input_shift: f64,
1134 derivative_stack: [f64; 5],
1135 &(): &'arena Self::Workspace,
1136 ) -> Self {
1137 Self {
1138 inner: self
1139 .inner
1140 .affine_compose(input_scale, input_shift, derivative_stack),
1141 }
1142 }
1143
1144 #[inline(always)]
1145 fn affine_composed_sum(
1146 inputs: &[Self],
1147 input_scales: &[f64],
1148 derivative_stacks: &[[f64; 5]],
1149 dimension: usize,
1150 &(): &'arena Self::Workspace,
1151 ) -> Self {
1152 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1153 let inner =
1157 unsafe { std::slice::from_raw_parts(inputs.as_ptr().cast::<S>(), inputs.len()) };
1158 Self {
1159 inner: S::affine_composed_sum(inner, input_scales, derivative_stacks),
1160 }
1161 }
1162
1163 #[inline(always)]
1164 fn shared_multiply_add_affine_composed_sum<const N: usize>(
1165 lefts: &[&Self; N],
1166 right: &Self,
1167 addend: &Self,
1168 addend_scales: &[f64; N],
1169 input_scales: &[f64; N],
1170 derivative_stacks: &[[f64; 5]; N],
1171 dimension: usize,
1172 &(): &'arena Self::Workspace,
1173 ) -> Self {
1174 assert_eq!(dimension, K, "fixed jet dimension mismatch");
1175 let left_inner: [&S; N] = std::array::from_fn(|term| &lefts[term].inner);
1176 Self {
1177 inner: S::shared_multiply_add_affine_composed_sum(
1178 &left_inner,
1179 &right.inner,
1180 &addend.inner,
1181 addend_scales,
1182 input_scales,
1183 derivative_stacks,
1184 ),
1185 }
1186 }
1187
1188 #[inline(always)]
1189 fn dimension(&self) -> usize {
1190 K
1191 }
1192
1193 #[inline(always)]
1194 fn value(&self) -> f64 {
1195 self.inner.value()
1196 }
1197
1198 #[inline(always)]
1199 fn add(&self, o: &Self) -> Self {
1200 Self {
1201 inner: self.inner.add(&o.inner),
1202 }
1203 }
1204
1205 #[inline(always)]
1206 fn sub(&self, o: &Self) -> Self {
1207 Self {
1208 inner: self.inner.sub(&o.inner),
1209 }
1210 }
1211
1212 #[inline(always)]
1213 fn mul(&self, o: &Self) -> Self {
1214 Self {
1215 inner: self.inner.mul(&o.inner),
1216 }
1217 }
1218
1219 #[inline(always)]
1220 fn neg(&self) -> Self {
1221 Self {
1222 inner: self.inner.neg(),
1223 }
1224 }
1225
1226 #[inline(always)]
1227 fn scale(&self, s: f64) -> Self {
1228 Self {
1229 inner: self.inner.scale(s),
1230 }
1231 }
1232
1233 #[inline(always)]
1234 fn compose_unary(&self, d: [f64; 5]) -> Self {
1235 Self {
1236 inner: self.inner.compose_unary(d),
1237 }
1238 }
1239}
1240
1241#[derive(Debug)]
1245pub struct DynamicJetArena {
1246 bump: bumpalo::Bump,
1247}
1248
1249impl DynamicJetArena {
1250 #[must_use]
1252 pub fn new() -> Self {
1253 Self {
1254 bump: bumpalo::Bump::new(),
1255 }
1256 }
1257
1258 #[must_use]
1260 pub fn with_capacity(bytes: usize) -> Self {
1261 Self {
1262 bump: bumpalo::Bump::with_capacity(bytes),
1263 }
1264 }
1265
1266 pub fn reset(&mut self) {
1277 let high_water = self.bump.allocated_bytes();
1278 self.bump.reset();
1279 if self.bump.allocated_bytes() < high_water {
1280 self.bump = bumpalo::Bump::with_capacity(high_water);
1281 }
1282 }
1283
1284 #[must_use]
1287 pub fn allocated_bytes(&self) -> usize {
1288 self.bump.allocated_bytes()
1289 }
1290
1291 #[inline(always)]
1292 fn zeros(&self, len: usize) -> &mut [f64] {
1293 self.bump.alloc_slice_fill_copy(len, 0.0)
1294 }
1295
1296 #[inline(always)]
1300 pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
1301 self.bump.alloc_slice_fill_with(len, fill)
1302 }
1303}
1304
1305impl Default for DynamicJetArena {
1306 fn default() -> Self {
1307 Self::new()
1308 }
1309}
1310
1311#[derive(Clone, Copy, Debug)]
1313pub struct DynamicOrder1<'arena> {
1314 arena: &'arena DynamicJetArena,
1315 pub v: f64,
1317 pub g: &'arena [f64],
1319}
1320
1321impl DynamicOrder1<'_> {
1322 #[inline]
1324 #[must_use]
1325 pub fn g(&self) -> &[f64] {
1326 self.g
1327 }
1328
1329 #[inline]
1330 fn assert_compatible(&self, o: &Self) {
1331 assert_eq!(
1332 self.g.len(),
1333 o.g.len(),
1334 "dynamic first-order jet dimension mismatch"
1335 );
1336 assert!(
1337 std::ptr::eq(self.arena, o.arena),
1338 "dynamic jets belong to different arenas"
1339 );
1340 }
1341}
1342
1343impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder1<'arena> {
1344 type Workspace = DynamicJetArena;
1345
1346 fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1347 Self {
1348 arena,
1349 v: c,
1350 g: arena.zeros(dimension),
1351 }
1352 }
1353
1354 fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1355 assert!(
1356 axis < dimension,
1357 "dynamic first-order jet axis out of bounds"
1358 );
1359 let g = arena.zeros(dimension);
1360 g[axis] = 1.0;
1361 Self { arena, v: x, g }
1362 }
1363
1364 fn dimension(&self) -> usize {
1365 self.g.len()
1366 }
1367 fn value(&self) -> f64 {
1368 self.v
1369 }
1370
1371 fn add(&self, o: &Self) -> Self {
1372 self.assert_compatible(o);
1373 let g = self.arena.zeros(self.dimension());
1374 for i in 0..g.len() {
1375 g[i] = self.g[i] + o.g[i];
1376 }
1377 Self {
1378 arena: self.arena,
1379 v: self.v + o.v,
1380 g,
1381 }
1382 }
1383
1384 fn sub(&self, o: &Self) -> Self {
1385 self.assert_compatible(o);
1386 let g = self.arena.zeros(self.dimension());
1387 for i in 0..g.len() {
1388 g[i] = self.g[i] - o.g[i];
1389 }
1390 Self {
1391 arena: self.arena,
1392 v: self.v - o.v,
1393 g,
1394 }
1395 }
1396
1397 fn mul(&self, o: &Self) -> Self {
1398 self.assert_compatible(o);
1399 let g = self.arena.zeros(self.dimension());
1400 for i in 0..g.len() {
1401 g[i] = self.v * o.g[i] + self.g[i] * o.v;
1402 }
1403 Self {
1404 arena: self.arena,
1405 v: self.v * o.v,
1406 g,
1407 }
1408 }
1409
1410 fn neg(&self) -> Self {
1411 self.scale(-1.0)
1412 }
1413
1414 fn scale(&self, s: f64) -> Self {
1415 let g = self.arena.zeros(self.dimension());
1416 for i in 0..g.len() {
1417 g[i] = self.g[i] * s;
1418 }
1419 Self {
1420 arena: self.arena,
1421 v: self.v * s,
1422 g,
1423 }
1424 }
1425
1426 fn compose_unary(&self, d: [f64; 5]) -> Self {
1427 let g = self.arena.zeros(self.dimension());
1428 for i in 0..g.len() {
1429 g[i] = d[1] * self.g[i];
1430 }
1431 Self {
1432 arena: self.arena,
1433 v: d[0],
1434 g,
1435 }
1436 }
1437}
1438
1439#[derive(Clone, Copy, Debug)]
1443pub struct DynamicOrder2<'arena> {
1444 arena: &'arena DynamicJetArena,
1445 pub v: f64,
1447 pub g: &'arena [f64],
1449 pub h: &'arena [f64],
1451}
1452
1453impl DynamicOrder2<'_> {
1454 #[inline]
1464 #[must_use]
1465 pub fn from_channel_functions<'arena>(
1466 value: f64,
1467 dimension: usize,
1468 arena: &'arena DynamicJetArena,
1469 mut gradient: impl FnMut(usize) -> f64,
1470 mut hessian: impl FnMut(usize, usize) -> f64,
1471 ) -> DynamicOrder2<'arena> {
1472 let g = arena.alloc_slice_fill_with(dimension, |axis| gradient(axis));
1473 let h = arena.zeros(dimension * dimension);
1474 for row in 0..dimension {
1475 for column in row..dimension {
1476 let channel = hessian(row, column);
1477 h[row * dimension + column] = channel;
1478 h[column * dimension + row] = channel;
1479 }
1480 }
1481 DynamicOrder2 {
1482 arena,
1483 v: value,
1484 g,
1485 h,
1486 }
1487 }
1488
1489 #[inline]
1491 #[must_use]
1492 pub fn g(&self) -> &[f64] {
1493 self.g
1494 }
1495
1496 #[inline]
1498 #[must_use]
1499 pub fn h(&self) -> &[f64] {
1500 self.h
1501 }
1502
1503 #[inline]
1505 #[must_use]
1506 pub fn h_at(&self, row: usize, col: usize) -> f64 {
1507 self.h[row * self.dimension() + col]
1508 }
1509
1510 #[inline(always)]
1511 fn assert_compatible(&self, o: &Self) {
1512 assert_eq!(
1513 self.g.len(),
1514 o.g.len(),
1515 "dynamic second-order jet dimension mismatch"
1516 );
1517 assert_eq!(
1518 self.h.len(),
1519 o.h.len(),
1520 "dynamic second-order jet Hessian mismatch"
1521 );
1522 assert!(
1523 std::ptr::eq(self.arena, o.arena),
1524 "dynamic jets belong to different arenas"
1525 );
1526 }
1527}
1528
1529impl<'arena> RuntimeJetScalar<'arena> for DynamicOrder2<'arena> {
1530 type Workspace = DynamicJetArena;
1531
1532 #[inline(always)]
1533 fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1534 Self {
1535 arena,
1536 v: c,
1537 g: arena.zeros(dimension),
1538 h: arena.zeros(dimension * dimension),
1539 }
1540 }
1541
1542 #[inline(always)]
1543 fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
1544 assert!(
1545 axis < dimension,
1546 "dynamic second-order jet axis out of bounds"
1547 );
1548 let g = arena.zeros(dimension);
1549 g[axis] = 1.0;
1550 Self {
1551 arena,
1552 v: x,
1553 g,
1554 h: arena.zeros(dimension * dimension),
1555 }
1556 }
1557
1558 #[inline(always)]
1559 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
1560 inputs: &[Self],
1561 coefficients: &C,
1562 dimension: usize,
1563 arena: &'arena DynamicJetArena,
1564 ) -> Self {
1565 assert_eq!(inputs.len(), coefficients.dimension());
1566 assert!(
1567 inputs.iter().all(|input| {
1568 input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1569 }),
1570 "dynamic quadratic-form jets must share dimension and arena"
1571 );
1572 let input_dimension = inputs.len();
1573 let values = arena.zeros(input_dimension);
1574 for (value, input) in values.iter_mut().zip(inputs) {
1575 *value = input.v;
1576 }
1577 let projected = arena.zeros(input_dimension);
1578 coefficients.multiply(values, projected);
1579
1580 let mut value = 0.0;
1581 for axis in 0..input_dimension {
1582 value += values[axis] * projected[axis];
1583 }
1584 let gradient = arena.zeros(dimension);
1585 for primary in 0..dimension {
1586 let mut channel = 0.0;
1587 for axis in 0..input_dimension {
1588 channel += projected[axis] * inputs[axis].g[primary];
1589 }
1590 gradient[primary] = 2.0 * channel;
1591 }
1592 let hessian = arena.zeros(dimension * dimension);
1593 let input_gradient = arena.zeros(input_dimension);
1594 let projected_gradient = arena.zeros(input_dimension);
1595 for primary_b in 0..dimension {
1596 for row in 0..input_dimension {
1597 input_gradient[row] = inputs[row].g[primary_b];
1598 }
1599 coefficients.multiply(input_gradient, projected_gradient);
1600 for primary_a in 0..=primary_b {
1601 let mut inherited = 0.0;
1602 let mut curvature = 0.0;
1603 for row in 0..input_dimension {
1604 inherited += projected[row] * inputs[row].h[primary_a * dimension + primary_b];
1605 curvature += inputs[row].g[primary_a] * projected_gradient[row];
1606 }
1607 let channel = 2.0 * (inherited + curvature);
1608 hessian[primary_a * dimension + primary_b] = channel;
1609 hessian[primary_b * dimension + primary_a] = channel;
1610 }
1611 }
1612 Self {
1613 arena,
1614 v: value,
1615 g: gradient,
1616 h: hessian,
1617 }
1618 }
1619
1620 #[inline(always)]
1621 fn product(&self, right: &Self) -> Self {
1622 self.assert_compatible(right);
1623 let dimension = self.dimension();
1624 let gradient = self.arena.zeros(dimension);
1625 let hessian = self.arena.zeros(dimension * dimension);
1626 for primary in 0..dimension {
1627 gradient[primary] = self.v * right.g[primary] + self.g[primary] * right.v;
1628 for other in primary..dimension {
1629 let index = primary * dimension + other;
1630 let channel = self.v * right.h[index]
1631 + self.g[primary] * right.g[other]
1632 + self.g[other] * right.g[primary]
1633 + self.h[index] * right.v;
1634 hessian[index] = channel;
1635 hessian[other * dimension + primary] = channel;
1636 }
1637 }
1638 Self {
1639 arena: self.arena,
1640 v: self.v * right.v,
1641 g: gradient,
1642 h: hessian,
1643 }
1644 }
1645
1646 #[inline(always)]
1647 fn affine_compose(
1648 &self,
1649 input_scale: f64,
1650 input_shift: f64,
1651 derivative_stack: [f64; 5],
1652 arena: &'arena DynamicJetArena,
1653 ) -> Self {
1654 assert!(std::ptr::eq(self.arena, arena));
1655 assert!(input_shift.is_finite(), "affine input shift must be finite");
1656 let dimension = self.dimension();
1657 let first = derivative_stack[1] * input_scale;
1658 let second = derivative_stack[2] * input_scale * input_scale;
1659 let gradient = arena.zeros(dimension);
1660 let hessian = arena.zeros(dimension * dimension);
1661 for primary in 0..dimension {
1662 gradient[primary] = first * self.g[primary];
1663 for other in primary..dimension {
1664 let index = primary * dimension + other;
1665 let channel = first * self.h[index] + second * self.g[primary] * self.g[other];
1666 hessian[index] = channel;
1667 hessian[other * dimension + primary] = channel;
1668 }
1669 }
1670 Self {
1671 arena,
1672 v: derivative_stack[0],
1673 g: gradient,
1674 h: hessian,
1675 }
1676 }
1677
1678 #[inline(always)]
1679 fn affine_composed_sum(
1680 inputs: &[Self],
1681 input_scales: &[f64],
1682 derivative_stacks: &[[f64; 5]],
1683 dimension: usize,
1684 arena: &'arena DynamicJetArena,
1685 ) -> Self {
1686 assert_eq!(inputs.len(), input_scales.len());
1687 assert_eq!(inputs.len(), derivative_stacks.len());
1688 assert!(
1689 inputs.iter().all(|input| {
1690 input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1691 }),
1692 "dynamic affine-composed-sum jets must share dimension and arena"
1693 );
1694 let gradient = arena.zeros(dimension);
1695 let hessian = arena.zeros(dimension * dimension);
1696 let mut value = 0.0;
1697 for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
1698 {
1699 let first = stack[1] * input_scale;
1700 let second = stack[2] * input_scale * input_scale;
1701 value += stack[0];
1702 for primary in 0..dimension {
1703 gradient[primary] += first * input.g[primary];
1704 for other in primary..dimension {
1705 let index = primary * dimension + other;
1706 hessian[index] +=
1707 first * input.h[index] + second * input.g[primary] * input.g[other];
1708 }
1709 }
1710 }
1711 for primary in 0..dimension {
1712 for other in primary + 1..dimension {
1713 hessian[other * dimension + primary] = hessian[primary * dimension + other];
1714 }
1715 }
1716 Self {
1717 arena,
1718 v: value,
1719 g: gradient,
1720 h: hessian,
1721 }
1722 }
1723
1724 #[inline(always)]
1725 fn shared_multiply_add_affine_composed_sum<const N: usize>(
1726 lefts: &[&Self; N],
1727 right: &Self,
1728 addend: &Self,
1729 addend_scales: &[f64; N],
1730 input_scales: &[f64; N],
1731 derivative_stacks: &[[f64; 5]; N],
1732 dimension: usize,
1733 arena: &'arena DynamicJetArena,
1734 ) -> Self {
1735 assert!(
1736 lefts.iter().all(|input| {
1737 input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1738 }) && (N == 0 || (right.dimension() == dimension && std::ptr::eq(right.arena, arena))),
1739 "dynamic fused product-composition jets must share dimension and arena"
1740 );
1741 let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
1742 assert!(
1743 !addend_live || (addend.dimension() == dimension && std::ptr::eq(addend.arena, arena)),
1744 "live dynamic fused addends must share dimension and arena"
1745 );
1746 let (representatives, term_sources, source_count) =
1747 canonical_shared_source_schedule::<N>(|term, representative| {
1748 std::ptr::eq(lefts[term], lefts[representative])
1749 && addend_scales[term] == addend_scales[representative]
1750 });
1751 let (value, source_derivatives) =
1752 aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
1753 let source_gradients = arena.zeros(source_count * dimension);
1754 let gradient = arena.zeros(dimension);
1755 let hessian = arena.zeros(dimension * dimension);
1756 let mut right_first = 0.0;
1757 let mut addend_first = 0.0;
1758 for source in 0..source_count {
1759 let term = representatives[source];
1760 let first = source_derivatives[source][1];
1761 right_first += first * lefts[term].v;
1762 addend_first += first * addend_scales[term];
1763 for primary in 0..dimension {
1764 let product_gradient =
1765 lefts[term].v * right.g[primary] + lefts[term].g[primary] * right.v;
1766 let inner_gradient = if addend_scales[term] == 0.0 {
1767 product_gradient
1768 } else if addend_scales[term] == 1.0 {
1769 product_gradient + addend.g[primary]
1770 } else {
1771 product_gradient + addend_scales[term] * addend.g[primary]
1772 };
1773 source_gradients[source * dimension + primary] = inner_gradient;
1774 gradient[primary] += first * lefts[term].g[primary] * right.v;
1775 }
1776 }
1777 if N != 0 {
1778 for primary in 0..dimension {
1779 gradient[primary] += right_first * right.g[primary];
1780 }
1781 }
1782 if addend_live {
1783 for primary in 0..dimension {
1784 gradient[primary] += addend_first * addend.g[primary];
1785 }
1786 }
1787 for primary in 0..dimension {
1788 for other in primary..dimension {
1789 let index = primary * dimension + other;
1790 let mut channel = if N == 0 {
1791 0.0
1792 } else {
1793 right_first * right.h[index]
1794 };
1795 if addend_live {
1796 channel += addend_first * addend.h[index];
1797 }
1798 for source in 0..source_count {
1799 let term = representatives[source];
1800 let local_product_hessian = lefts[term].g[primary] * right.g[other]
1801 + lefts[term].g[other] * right.g[primary]
1802 + lefts[term].h[index] * right.v;
1803 channel += source_derivatives[source][1] * local_product_hessian
1804 + source_derivatives[source][2]
1805 * source_gradients[source * dimension + primary]
1806 * source_gradients[source * dimension + other];
1807 }
1808 hessian[index] = channel;
1809 hessian[other * dimension + primary] = channel;
1810 }
1811 }
1812 Self {
1813 arena,
1814 v: value,
1815 g: gradient,
1816 h: hessian,
1817 }
1818 }
1819
1820 #[inline(always)]
1821 fn add_constant(&self, constant: f64, arena: &'arena DynamicJetArena) -> Self {
1822 assert!(std::ptr::eq(self.arena, arena));
1823 Self {
1824 arena,
1825 v: self.v + constant,
1826 g: self.g,
1827 h: self.h,
1828 }
1829 }
1830
1831 #[inline(always)]
1832 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
1833 self.assert_compatible(right);
1834 self.assert_compatible(addend);
1835 let dimension = self.dimension();
1836 let gradient = self.arena.zeros(dimension);
1837 let hessian = self.arena.zeros(dimension * dimension);
1838 for primary in 0..dimension {
1839 gradient[primary] =
1840 self.v * right.g[primary] + self.g[primary] * right.v + addend.g[primary];
1841 for other in primary..dimension {
1842 let index = primary * dimension + other;
1843 let channel = self.v * right.h[index]
1844 + self.g[primary] * right.g[other]
1845 + self.g[other] * right.g[primary]
1846 + self.h[index] * right.v
1847 + addend.h[index];
1848 hessian[index] = channel;
1849 hessian[other * dimension + primary] = channel;
1850 }
1851 }
1852 Self {
1853 arena: self.arena,
1854 v: self.v * right.v + addend.v,
1855 g: gradient,
1856 h: hessian,
1857 }
1858 }
1859
1860 #[inline(always)]
1861 fn composed_sum(
1862 inputs: &[Self],
1863 derivative_stacks: &[[f64; 5]],
1864 dimension: usize,
1865 arena: &'arena DynamicJetArena,
1866 ) -> Self {
1867 assert_eq!(inputs.len(), derivative_stacks.len());
1868 assert!(
1869 inputs.iter().all(|input| {
1870 input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1871 }),
1872 "dynamic composed-sum jets must share dimension and arena"
1873 );
1874 let gradient = arena.zeros(dimension);
1875 let hessian = arena.zeros(dimension * dimension);
1876 let mut value = 0.0;
1877 for (input, stack) in inputs.iter().zip(derivative_stacks) {
1878 value += stack[0];
1879 for primary in 0..dimension {
1880 gradient[primary] += stack[1] * input.g[primary];
1881 for other in primary..dimension {
1882 let index = primary * dimension + other;
1883 hessian[index] +=
1884 stack[1] * input.h[index] + stack[2] * input.g[primary] * input.g[other];
1885 }
1886 }
1887 }
1888 for primary in 0..dimension {
1889 for other in primary + 1..dimension {
1890 hessian[other * dimension + primary] = hessian[primary * dimension + other];
1891 }
1892 }
1893 Self {
1894 arena,
1895 v: value,
1896 g: gradient,
1897 h: hessian,
1898 }
1899 }
1900
1901 #[inline(always)]
1902 fn linear_combination(
1903 inputs: &[Self],
1904 weights: &[f64],
1905 dimension: usize,
1906 arena: &'arena DynamicJetArena,
1907 ) -> Self {
1908 assert_eq!(inputs.len(), weights.len());
1909 assert!(
1910 inputs.iter().all(|input| {
1911 input.dimension() == dimension && std::ptr::eq(input.arena, arena)
1912 }),
1913 "dynamic linear-combination jets must share dimension and arena"
1914 );
1915 let mut value = 0.0;
1916 for (input, &weight) in inputs.iter().zip(weights) {
1917 value += input.v * weight;
1918 }
1919 let gradient = arena.zeros(dimension);
1920 let hessian = arena.zeros(dimension * dimension);
1921 for primary in 0..dimension {
1922 for (input, &weight) in inputs.iter().zip(weights) {
1923 gradient[primary] += input.g[primary] * weight;
1924 }
1925 for other in primary..dimension {
1926 let index = primary * dimension + other;
1927 for (input, &weight) in inputs.iter().zip(weights) {
1928 hessian[index] += input.h[index] * weight;
1929 }
1930 hessian[other * dimension + primary] = hessian[index];
1931 }
1932 }
1933 Self {
1934 arena,
1935 v: value,
1936 g: gradient,
1937 h: hessian,
1938 }
1939 }
1940
1941 #[inline(always)]
1942 fn dimension(&self) -> usize {
1943 self.g.len()
1944 }
1945
1946 #[inline(always)]
1947 fn value(&self) -> f64 {
1948 self.v
1949 }
1950
1951 #[inline(always)]
1952 fn add(&self, o: &Self) -> Self {
1953 self.assert_compatible(o);
1954 let dimension = self.dimension();
1955 let g = self.arena.zeros(dimension);
1956 let h = self.arena.zeros(self.h.len());
1957 for i in 0..g.len() {
1958 g[i] = self.g[i] + o.g[i];
1959 }
1960 for row in 0..dimension {
1961 for column in row..dimension {
1962 let index = row * dimension + column;
1963 let channel = self.h[index] + o.h[index];
1964 h[index] = channel;
1965 h[column * dimension + row] = channel;
1966 }
1967 }
1968 Self {
1969 arena: self.arena,
1970 v: self.v + o.v,
1971 g,
1972 h,
1973 }
1974 }
1975
1976 #[inline(always)]
1977 fn sub(&self, o: &Self) -> Self {
1978 self.assert_compatible(o);
1979 let dimension = self.dimension();
1980 let g = self.arena.zeros(dimension);
1981 let h = self.arena.zeros(self.h.len());
1982 for i in 0..g.len() {
1983 g[i] = self.g[i] - o.g[i];
1984 }
1985 for row in 0..dimension {
1986 for column in row..dimension {
1987 let index = row * dimension + column;
1988 let channel = self.h[index] - o.h[index];
1989 h[index] = channel;
1990 h[column * dimension + row] = channel;
1991 }
1992 }
1993 Self {
1994 arena: self.arena,
1995 v: self.v - o.v,
1996 g,
1997 h,
1998 }
1999 }
2000
2001 #[inline(always)]
2002 fn mul(&self, o: &Self) -> Self {
2003 self.assert_compatible(o);
2004 let n = self.dimension();
2005 let g = self.arena.zeros(n);
2006 let h = self.arena.zeros(n * n);
2007 for i in 0..n {
2008 g[i] = self.v * o.g[i] + self.g[i] * o.v;
2009 }
2010 for i in 0..n {
2011 for j in i..n {
2012 let ij = i * n + j;
2013 let hij =
2014 self.v * o.h[ij] + self.g[i] * o.g[j] + self.g[j] * o.g[i] + self.h[ij] * o.v;
2015 h[ij] = hij;
2016 h[j * n + i] = hij;
2017 }
2018 }
2019 Self {
2020 arena: self.arena,
2021 v: self.v * o.v,
2022 g,
2023 h,
2024 }
2025 }
2026
2027 #[inline(always)]
2028 fn neg(&self) -> Self {
2029 self.scale(-1.0)
2030 }
2031
2032 #[inline(always)]
2033 fn scale(&self, s: f64) -> Self {
2034 let dimension = self.dimension();
2035 let g = self.arena.zeros(dimension);
2036 let h = self.arena.zeros(self.h.len());
2037 for i in 0..g.len() {
2038 g[i] = self.g[i] * s;
2039 }
2040 for row in 0..dimension {
2041 for column in row..dimension {
2042 let index = row * dimension + column;
2043 let channel = self.h[index] * s;
2044 h[index] = channel;
2045 h[column * dimension + row] = channel;
2046 }
2047 }
2048 Self {
2049 arena: self.arena,
2050 v: self.v * s,
2051 g,
2052 h,
2053 }
2054 }
2055
2056 #[inline(always)]
2057 fn compose_unary(&self, d: [f64; 5]) -> Self {
2058 let n = self.dimension();
2059 let g = self.arena.zeros(n);
2060 let h = self.arena.zeros(n * n);
2061 for i in 0..n {
2062 g[i] = d[1] * self.g[i];
2063 }
2064 for i in 0..n {
2065 for j in i..n {
2066 let ij = i * n + j;
2067 let channel = d[1] * self.h[ij] + d[2] * self.g[i] * self.g[j];
2068 h[ij] = channel;
2069 h[j * n + i] = channel;
2070 }
2071 }
2072 Self {
2073 arena: self.arena,
2074 v: d[0],
2075 g,
2076 h,
2077 }
2078 }
2079}
2080
2081#[derive(Clone, Copy, Debug)]
2083pub struct DynamicOneSeed<'arena> {
2084 pub base: DynamicOrder2<'arena>,
2086 pub eps: DynamicOrder2<'arena>,
2088}
2089
2090impl<'arena> DynamicOneSeed<'arena> {
2091 #[inline(always)]
2093 #[must_use]
2094 pub fn seed_direction(
2095 x: f64,
2096 axis: usize,
2097 u_axis: f64,
2098 dimension: usize,
2099 arena: &'arena DynamicJetArena,
2100 ) -> Self {
2101 Self {
2102 base: DynamicOrder2::variable(x, axis, dimension, arena),
2103 eps: DynamicOrder2::constant(u_axis, dimension, arena),
2104 }
2105 }
2106
2107 #[inline(always)]
2109 #[must_use]
2110 pub fn contracted_third(&self) -> &[f64] {
2111 self.eps.h()
2112 }
2113}
2114
2115impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeed<'arena> {
2116 type Workspace = DynamicJetArena;
2117
2118 #[inline(always)]
2119 fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2120 Self {
2121 base: DynamicOrder2::constant(c, dimension, arena),
2122 eps: DynamicOrder2::constant(0.0, dimension, arena),
2123 }
2124 }
2125
2126 #[inline(always)]
2127 fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2128 Self {
2129 base: DynamicOrder2::variable(x, axis, dimension, arena),
2130 eps: DynamicOrder2::constant(0.0, dimension, arena),
2131 }
2132 }
2133
2134 #[inline(always)]
2135 fn dimension(&self) -> usize {
2136 self.base.dimension()
2137 }
2138
2139 #[inline(always)]
2140 fn value(&self) -> f64 {
2141 self.base.value()
2142 }
2143
2144 #[inline(always)]
2145 fn add(&self, o: &Self) -> Self {
2146 Self {
2147 base: self.base.add(&o.base),
2148 eps: self.eps.add(&o.eps),
2149 }
2150 }
2151
2152 #[inline(always)]
2153 fn sub(&self, o: &Self) -> Self {
2154 Self {
2155 base: self.base.sub(&o.base),
2156 eps: self.eps.sub(&o.eps),
2157 }
2158 }
2159
2160 #[inline(always)]
2161 fn mul(&self, o: &Self) -> Self {
2162 self.base.assert_compatible(&o.base);
2163 self.eps.assert_compatible(&o.eps);
2164 Self {
2165 base: self.base.mul(&o.base),
2166 eps: DynamicOrder2::from_channel_functions(
2167 self.base.v * o.eps.v + self.eps.v * o.base.v,
2168 self.dimension(),
2169 self.base.arena,
2170 |i| {
2171 self.base.v * o.eps.g[i]
2172 + self.base.g[i] * o.eps.v
2173 + self.eps.v * o.base.g[i]
2174 + self.eps.g[i] * o.base.v
2175 },
2176 |i, j| {
2177 let ij = i * self.dimension() + j;
2178 self.base.v * o.eps.h[ij]
2179 + self.base.g[i] * o.eps.g[j]
2180 + self.base.g[j] * o.eps.g[i]
2181 + self.base.h[ij] * o.eps.v
2182 + self.eps.v * o.base.h[ij]
2183 + self.eps.g[i] * o.base.g[j]
2184 + self.eps.g[j] * o.base.g[i]
2185 + self.eps.h[ij] * o.base.v
2186 },
2187 ),
2188 }
2189 }
2190
2191 #[inline(always)]
2192 fn neg(&self) -> Self {
2193 Self {
2194 base: self.base.neg(),
2195 eps: self.eps.neg(),
2196 }
2197 }
2198
2199 #[inline(always)]
2200 fn scale(&self, s: f64) -> Self {
2201 Self {
2202 base: self.base.scale(s),
2203 eps: self.eps.scale(s),
2204 }
2205 }
2206
2207 #[inline(always)]
2208 fn compose_unary(&self, d: [f64; 5]) -> Self {
2209 let base = self.base.compose_unary(d);
2210 let dimension = self.dimension();
2211 let eps = DynamicOrder2::from_channel_functions(
2212 d[1] * self.eps.v,
2213 dimension,
2214 self.base.arena,
2215 |i| d[2] * self.base.g[i] * self.eps.v + d[1] * self.eps.g[i],
2216 |i, j| {
2217 let ij = i * dimension + j;
2218 d[1] * self.eps.h[ij]
2219 + d[2]
2220 * (self.base.g[i] * self.eps.g[j]
2221 + self.base.g[j] * self.eps.g[i]
2222 + self.base.h[ij] * self.eps.v)
2223 + d[3] * self.base.g[i] * self.base.g[j] * self.eps.v
2224 },
2225 );
2226 Self { base, eps }
2227 }
2228}
2229
2230#[derive(Debug)]
2237pub struct DynamicJetBatchWorkspace {
2238 arena: DynamicJetArena,
2239 lanes: usize,
2240}
2241
2242impl DynamicJetBatchWorkspace {
2243 #[must_use]
2245 pub fn new(lanes: usize) -> Self {
2246 Self {
2247 arena: DynamicJetArena::new(),
2248 lanes,
2249 }
2250 }
2251
2252 pub fn reset(&mut self, lanes: usize) {
2254 self.arena.reset();
2255 self.lanes = lanes;
2256 }
2257
2258 #[must_use]
2260 pub fn allocated_bytes(&self) -> usize {
2261 self.arena.allocated_bytes()
2262 }
2263
2264 #[inline(always)]
2266 pub fn alloc_slice_fill_with<T>(&self, len: usize, fill: impl FnMut(usize) -> T) -> &mut [T] {
2267 self.arena.alloc_slice_fill_with(len, fill)
2268 }
2269}
2270
2271#[derive(Clone, Copy, Debug)]
2278pub struct DynamicOneSeedBatch<'arena> {
2279 pub base: DynamicOrder2<'arena>,
2281 eps: &'arena [DynamicOrder2<'arena>],
2283}
2284
2285impl<'arena> DynamicOneSeedBatch<'arena> {
2286 #[inline(always)]
2288 #[must_use]
2289 pub fn seed_directions(
2290 x: f64,
2291 axis: usize,
2292 dimension: usize,
2293 workspace: &'arena DynamicJetBatchWorkspace,
2294 mut direction_at: impl FnMut(usize) -> f64,
2295 ) -> Self {
2296 let eps = workspace
2297 .arena
2298 .alloc_slice_fill_with(workspace.lanes, |lane| {
2299 DynamicOrder2::constant(direction_at(lane), dimension, &workspace.arena)
2300 });
2301 Self {
2302 base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2303 eps,
2304 }
2305 }
2306
2307 #[inline(always)]
2309 #[must_use]
2310 pub fn lanes(&self) -> usize {
2311 self.eps.len()
2312 }
2313
2314 #[inline(always)]
2316 #[must_use]
2317 pub fn contracted_third(&self, lane: usize) -> &[f64] {
2318 self.eps[lane].h()
2319 }
2320
2321 #[inline(always)]
2322 fn assert_compatible(&self, other: &Self) {
2323 self.base.assert_compatible(&other.base);
2324 assert_eq!(
2325 self.eps.len(),
2326 other.eps.len(),
2327 "dynamic one-seed batch lane mismatch"
2328 );
2329 }
2330}
2331
2332impl<'arena> RuntimeJetScalar<'arena> for DynamicOneSeedBatch<'arena> {
2333 type Workspace = DynamicJetBatchWorkspace;
2334
2335 #[inline(always)]
2336 fn constant(c: f64, dimension: usize, workspace: &'arena DynamicJetBatchWorkspace) -> Self {
2337 let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2338 DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2339 });
2340 Self {
2341 base: DynamicOrder2::constant(c, dimension, &workspace.arena),
2342 eps,
2343 }
2344 }
2345
2346 #[inline(always)]
2347 fn variable(
2348 x: f64,
2349 axis: usize,
2350 dimension: usize,
2351 workspace: &'arena DynamicJetBatchWorkspace,
2352 ) -> Self {
2353 let eps = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2354 DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2355 });
2356 Self {
2357 base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2358 eps,
2359 }
2360 }
2361
2362 #[inline(always)]
2363 fn dimension(&self) -> usize {
2364 self.base.dimension()
2365 }
2366
2367 #[inline(always)]
2368 fn value(&self) -> f64 {
2369 self.base.value()
2370 }
2371
2372 #[inline(always)]
2373 fn add(&self, other: &Self) -> Self {
2374 self.assert_compatible(other);
2375 let eps = self
2376 .base
2377 .arena
2378 .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].add(&other.eps[lane]));
2379 Self {
2380 base: self.base.add(&other.base),
2381 eps,
2382 }
2383 }
2384
2385 #[inline(always)]
2386 fn sub(&self, other: &Self) -> Self {
2387 self.assert_compatible(other);
2388 let eps = self
2389 .base
2390 .arena
2391 .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].sub(&other.eps[lane]));
2392 Self {
2393 base: self.base.sub(&other.base),
2394 eps,
2395 }
2396 }
2397
2398 #[inline(always)]
2399 fn mul(&self, other: &Self) -> Self {
2400 self.assert_compatible(other);
2401 let eps = self
2402 .base
2403 .arena
2404 .alloc_slice_fill_with(self.eps.len(), |lane| {
2405 self.base
2406 .mul(&other.eps[lane])
2407 .add(&self.eps[lane].mul(&other.base))
2408 });
2409 Self {
2410 base: self.base.mul(&other.base),
2411 eps,
2412 }
2413 }
2414
2415 #[inline(always)]
2416 fn neg(&self) -> Self {
2417 self.scale(-1.0)
2418 }
2419
2420 #[inline(always)]
2421 fn scale(&self, scale: f64) -> Self {
2422 let eps = self
2423 .base
2424 .arena
2425 .alloc_slice_fill_with(self.eps.len(), |lane| self.eps[lane].scale(scale));
2426 Self {
2427 base: self.base.scale(scale),
2428 eps,
2429 }
2430 }
2431
2432 #[inline(always)]
2433 fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
2434 let fprime = self.base.compose_unary([
2435 derivatives[1],
2436 derivatives[2],
2437 derivatives[3],
2438 derivatives[4],
2439 derivatives[4],
2440 ]);
2441 let eps = self
2442 .base
2443 .arena
2444 .alloc_slice_fill_with(self.eps.len(), |lane| fprime.mul(&self.eps[lane]));
2445 Self {
2446 base: self.base.compose_unary(derivatives),
2447 eps,
2448 }
2449 }
2450}
2451
2452#[derive(Clone, Copy, Debug)]
2459pub struct DynamicTwoSeedBatch<'arena> {
2460 pub base: DynamicOrder2<'arena>,
2462 eps: &'arena [DynamicOrder2<'arena>],
2463 del: &'arena [DynamicOrder2<'arena>],
2464 eps_del: &'arena [DynamicOrder2<'arena>],
2465}
2466
2467impl<'arena> DynamicTwoSeedBatch<'arena> {
2468 #[inline(always)]
2470 #[must_use]
2471 pub fn seed_direction_pairs(
2472 x: f64,
2473 axis: usize,
2474 dimension: usize,
2475 workspace: &'arena DynamicJetBatchWorkspace,
2476 mut direction_pair_at: impl FnMut(usize) -> (f64, f64),
2477 ) -> Self {
2478 let directions = workspace
2479 .arena
2480 .alloc_slice_fill_with(workspace.lanes, |lane| direction_pair_at(lane));
2481 let eps = workspace
2482 .arena
2483 .alloc_slice_fill_with(workspace.lanes, |lane| {
2484 DynamicOrder2::constant(directions[lane].0, dimension, &workspace.arena)
2485 });
2486 let del = workspace
2487 .arena
2488 .alloc_slice_fill_with(workspace.lanes, |lane| {
2489 DynamicOrder2::constant(directions[lane].1, dimension, &workspace.arena)
2490 });
2491 let eps_del = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2492 DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2493 });
2494 Self {
2495 base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2496 eps,
2497 del,
2498 eps_del,
2499 }
2500 }
2501
2502 #[inline(always)]
2504 #[must_use]
2505 pub fn lanes(&self) -> usize {
2506 self.eps.len()
2507 }
2508
2509 #[inline(always)]
2511 #[must_use]
2512 pub fn contracted_fourth(&self, lane: usize) -> &[f64] {
2513 self.eps_del[lane].h()
2514 }
2515
2516 #[inline(always)]
2517 fn assert_compatible(&self, other: &Self) {
2518 self.base.assert_compatible(&other.base);
2519 assert_eq!(
2520 self.eps.len(),
2521 other.eps.len(),
2522 "dynamic two-seed batch lane mismatch"
2523 );
2524 assert_eq!(
2525 self.del.len(),
2526 self.eps.len(),
2527 "dynamic two-seed batch delta mismatch"
2528 );
2529 assert_eq!(
2530 self.eps_del.len(),
2531 self.eps.len(),
2532 "dynamic two-seed batch cross mismatch"
2533 );
2534 }
2535}
2536
2537impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeedBatch<'arena> {
2538 type Workspace = DynamicJetBatchWorkspace;
2539
2540 #[inline(always)]
2541 fn constant(c: f64, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
2542 let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2543 DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2544 });
2545 Self {
2546 base: DynamicOrder2::constant(c, dimension, &workspace.arena),
2547 eps: zero,
2548 del: zero,
2549 eps_del: zero,
2550 }
2551 }
2552
2553 #[inline(always)]
2554 fn variable(x: f64, axis: usize, dimension: usize, workspace: &'arena Self::Workspace) -> Self {
2555 let zero = workspace.arena.alloc_slice_fill_with(workspace.lanes, |_| {
2556 DynamicOrder2::constant(0.0, dimension, &workspace.arena)
2557 });
2558 Self {
2559 base: DynamicOrder2::variable(x, axis, dimension, &workspace.arena),
2560 eps: zero,
2561 del: zero,
2562 eps_del: zero,
2563 }
2564 }
2565
2566 #[inline(always)]
2567 fn dimension(&self) -> usize {
2568 self.base.dimension()
2569 }
2570
2571 #[inline(always)]
2572 fn value(&self) -> f64 {
2573 self.base.value()
2574 }
2575
2576 #[inline(always)]
2577 fn add(&self, other: &Self) -> Self {
2578 self.assert_compatible(other);
2579 let arena = self.base.arena;
2580 let eps =
2581 arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].add(&other.eps[lane]));
2582 let del =
2583 arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].add(&other.del[lane]));
2584 let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2585 self.eps_del[lane].add(&other.eps_del[lane])
2586 });
2587 Self {
2588 base: self.base.add(&other.base),
2589 eps,
2590 del,
2591 eps_del,
2592 }
2593 }
2594
2595 #[inline(always)]
2596 fn sub(&self, other: &Self) -> Self {
2597 self.assert_compatible(other);
2598 let arena = self.base.arena;
2599 let eps =
2600 arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].sub(&other.eps[lane]));
2601 let del =
2602 arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].sub(&other.del[lane]));
2603 let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2604 self.eps_del[lane].sub(&other.eps_del[lane])
2605 });
2606 Self {
2607 base: self.base.sub(&other.base),
2608 eps,
2609 del,
2610 eps_del,
2611 }
2612 }
2613
2614 #[inline(always)]
2615 fn mul(&self, other: &Self) -> Self {
2616 self.assert_compatible(other);
2617 let arena = self.base.arena;
2618 let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2619 self.base
2620 .mul(&other.eps[lane])
2621 .add(&self.eps[lane].mul(&other.base))
2622 });
2623 let del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2624 self.base
2625 .mul(&other.del[lane])
2626 .add(&self.del[lane].mul(&other.base))
2627 });
2628 let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2629 self.base
2630 .mul(&other.eps_del[lane])
2631 .add(&self.eps[lane].mul(&other.del[lane]))
2632 .add(&self.del[lane].mul(&other.eps[lane]))
2633 .add(&self.eps_del[lane].mul(&other.base))
2634 });
2635 Self {
2636 base: self.base.mul(&other.base),
2637 eps,
2638 del,
2639 eps_del,
2640 }
2641 }
2642
2643 #[inline(always)]
2644 fn neg(&self) -> Self {
2645 self.scale(-1.0)
2646 }
2647
2648 #[inline(always)]
2649 fn scale(&self, scale: f64) -> Self {
2650 let arena = self.base.arena;
2651 let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps[lane].scale(scale));
2652 let del = arena.alloc_slice_fill_with(self.lanes(), |lane| self.del[lane].scale(scale));
2653 let eps_del =
2654 arena.alloc_slice_fill_with(self.lanes(), |lane| self.eps_del[lane].scale(scale));
2655 Self {
2656 base: self.base.scale(scale),
2657 eps,
2658 del,
2659 eps_del,
2660 }
2661 }
2662
2663 #[inline(always)]
2664 fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
2665 let arena = self.base.arena;
2666 let fprime = self.base.compose_unary([
2667 derivatives[1],
2668 derivatives[2],
2669 derivatives[3],
2670 derivatives[4],
2671 derivatives[4],
2672 ]);
2673 let fsecond = self.base.compose_unary([
2674 derivatives[2],
2675 derivatives[3],
2676 derivatives[4],
2677 derivatives[4],
2678 derivatives[4],
2679 ]);
2680 let eps = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.eps[lane]));
2681 let del = arena.alloc_slice_fill_with(self.lanes(), |lane| fprime.mul(&self.del[lane]));
2682 let eps_del = arena.alloc_slice_fill_with(self.lanes(), |lane| {
2683 fsecond
2684 .mul(&self.eps[lane])
2685 .mul(&self.del[lane])
2686 .add(&fprime.mul(&self.eps_del[lane]))
2687 });
2688 Self {
2689 base: self.base.compose_unary(derivatives),
2690 eps,
2691 del,
2692 eps_del,
2693 }
2694 }
2695}
2696
2697#[derive(Clone, Copy, Debug)]
2699pub struct DynamicTwoSeed<'arena> {
2700 pub base: DynamicOrder2<'arena>,
2702 pub eps: DynamicOrder2<'arena>,
2704 pub del: DynamicOrder2<'arena>,
2706 pub eps_del: DynamicOrder2<'arena>,
2708}
2709
2710impl<'arena> DynamicTwoSeed<'arena> {
2711 #[inline(always)]
2713 #[must_use]
2714 pub fn seed(
2715 x: f64,
2716 axis: usize,
2717 u_axis: f64,
2718 v_axis: f64,
2719 dimension: usize,
2720 arena: &'arena DynamicJetArena,
2721 ) -> Self {
2722 Self {
2723 base: DynamicOrder2::variable(x, axis, dimension, arena),
2724 eps: DynamicOrder2::constant(u_axis, dimension, arena),
2725 del: DynamicOrder2::constant(v_axis, dimension, arena),
2726 eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2727 }
2728 }
2729
2730 #[inline(always)]
2732 #[must_use]
2733 pub fn contracted_fourth(&self) -> &[f64] {
2734 self.eps_del.h()
2735 }
2736}
2737
2738impl<'arena> RuntimeJetScalar<'arena> for DynamicTwoSeed<'arena> {
2739 type Workspace = DynamicJetArena;
2740
2741 #[inline(always)]
2742 fn constant(c: f64, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2743 Self {
2744 base: DynamicOrder2::constant(c, dimension, arena),
2745 eps: DynamicOrder2::constant(0.0, dimension, arena),
2746 del: DynamicOrder2::constant(0.0, dimension, arena),
2747 eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2748 }
2749 }
2750
2751 #[inline(always)]
2752 fn variable(x: f64, axis: usize, dimension: usize, arena: &'arena DynamicJetArena) -> Self {
2753 Self {
2754 base: DynamicOrder2::variable(x, axis, dimension, arena),
2755 eps: DynamicOrder2::constant(0.0, dimension, arena),
2756 del: DynamicOrder2::constant(0.0, dimension, arena),
2757 eps_del: DynamicOrder2::constant(0.0, dimension, arena),
2758 }
2759 }
2760
2761 #[inline(always)]
2762 fn dimension(&self) -> usize {
2763 self.base.dimension()
2764 }
2765
2766 #[inline(always)]
2767 fn value(&self) -> f64 {
2768 self.base.value()
2769 }
2770
2771 #[inline(always)]
2772 fn add(&self, o: &Self) -> Self {
2773 Self {
2774 base: self.base.add(&o.base),
2775 eps: self.eps.add(&o.eps),
2776 del: self.del.add(&o.del),
2777 eps_del: self.eps_del.add(&o.eps_del),
2778 }
2779 }
2780
2781 #[inline(always)]
2782 fn sub(&self, o: &Self) -> Self {
2783 Self {
2784 base: self.base.sub(&o.base),
2785 eps: self.eps.sub(&o.eps),
2786 del: self.del.sub(&o.del),
2787 eps_del: self.eps_del.sub(&o.eps_del),
2788 }
2789 }
2790
2791 #[inline(always)]
2792 fn mul(&self, o: &Self) -> Self {
2793 let base = self.base.mul(&o.base);
2794 let eps = self.base.mul(&o.eps).add(&self.eps.mul(&o.base));
2795 let del = self.base.mul(&o.del).add(&self.del.mul(&o.base));
2796 let eps_del = self
2797 .base
2798 .mul(&o.eps_del)
2799 .add(&self.eps.mul(&o.del))
2800 .add(&self.del.mul(&o.eps))
2801 .add(&self.eps_del.mul(&o.base));
2802 Self {
2803 base,
2804 eps,
2805 del,
2806 eps_del,
2807 }
2808 }
2809
2810 #[inline(always)]
2811 fn neg(&self) -> Self {
2812 Self {
2813 base: self.base.neg(),
2814 eps: self.eps.neg(),
2815 del: self.del.neg(),
2816 eps_del: self.eps_del.neg(),
2817 }
2818 }
2819
2820 #[inline(always)]
2821 fn scale(&self, s: f64) -> Self {
2822 Self {
2823 base: self.base.scale(s),
2824 eps: self.eps.scale(s),
2825 del: self.del.scale(s),
2826 eps_del: self.eps_del.scale(s),
2827 }
2828 }
2829
2830 #[inline(always)]
2831 fn compose_unary(&self, d: [f64; 5]) -> Self {
2832 let base = self.base.compose_unary(d);
2833 let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]);
2834 let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]);
2835 let eps = fprime.mul(&self.eps);
2836 let del = fprime.mul(&self.del);
2837 let eps_del = fsecond
2838 .mul(&self.eps)
2839 .mul(&self.del)
2840 .add(&fprime.mul(&self.eps_del));
2841 Self {
2842 base,
2843 eps,
2844 del,
2845 eps_del,
2846 }
2847 }
2848}
2849
2850impl<const K: usize> std::ops::Add for Order2<K> {
2860 type Output = Self;
2861 #[inline]
2862 fn add(self, o: Self) -> Self {
2863 Order2(self.0 + o.0)
2864 }
2865}
2866
2867impl<const K: usize> std::ops::Add<f64> for Order2<K> {
2868 type Output = Self;
2869 #[inline]
2870 fn add(self, c: f64) -> Self {
2871 Order2(self.0 + c)
2872 }
2873}
2874
2875impl<const K: usize> std::ops::Sub for Order2<K> {
2876 type Output = Self;
2877 #[inline]
2878 fn sub(self, o: Self) -> Self {
2879 Order2(self.0 + o.0.scale(-1.0))
2880 }
2881}
2882
2883impl<const K: usize> std::ops::Sub<f64> for Order2<K> {
2884 type Output = Self;
2885 #[inline]
2886 fn sub(self, c: f64) -> Self {
2887 Order2(self.0 + (-c))
2888 }
2889}
2890
2891impl<const K: usize> std::ops::Mul for Order2<K> {
2892 type Output = Self;
2893 #[inline]
2894 fn mul(self, o: Self) -> Self {
2895 Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
2896 }
2897}
2898
2899impl<const K: usize> std::ops::Mul<f64> for Order2<K> {
2900 type Output = Self;
2901 #[inline]
2902 fn mul(self, c: f64) -> Self {
2903 Order2(self.0.scale(c))
2904 }
2905}
2906
2907impl<const K: usize> std::ops::Neg for Order2<K> {
2908 type Output = Self;
2909 #[inline]
2910 fn neg(self) -> Self {
2911 Order2(self.0.scale(-1.0))
2912 }
2913}
2914
2915pub fn filtered_implicit_solve_scalar<const K: usize, S: JetScalar<K>>(
2940 a0: f64,
2941 inv_fa: f64,
2942 iters: usize,
2943 f: impl Fn(&S) -> S,
2944) -> S {
2945 let mut a = S::constant(a0);
2946 for _ in 0..iters {
2947 let residual = f(&a);
2948 a = a.sub(&residual.scale(inv_fa));
2949 }
2950 a
2951}
2952
2953pub fn filtered_implicit_solve_runtime_scalar<'arena, S: RuntimeJetScalar<'arena>>(
2965 a0: f64,
2966 inv_fa: f64,
2967 iters: usize,
2968 dimension: usize,
2969 workspace: &'arena S::Workspace,
2970 f: impl Fn(&S) -> S,
2971) -> S {
2972 let mut a = S::constant(a0, dimension, workspace);
2973 for _ in 0..iters {
2974 let residual = f(&a);
2975 a = a.sub(&residual.scale(inv_fa));
2976 }
2977 a
2978}
2979
2980pub trait HessianPattern<const K: usize, const H: usize> {
2988 const PAIRS: [(usize, usize); H];
2989 const PAIR_BITS: [[u128; K]; K];
2990}
2991
2992pub const fn hessian_pair_bits<const K: usize, const H: usize>(
2995 pairs: [(usize, usize); H],
2996) -> [[u128; K]; K] {
2997 let mut table = [[0u128; K]; K];
2998 let mut slot = 0;
2999 while slot < H {
3000 let (i, j) = pairs[slot];
3001 let bit = 1u128 << slot;
3002 table[i][j] = bit;
3003 table[j][i] = bit;
3004 slot += 1;
3005 }
3006 table
3007}
3008
3009#[derive(Debug)]
3018pub struct PatternedOrder2<P, const K: usize, const H: usize> {
3019 v: f64,
3020 g: [f64; K],
3021 h: [f64; H],
3022 gradient_mask: u128,
3023 hessian_mask: u128,
3024 pattern: std::marker::PhantomData<fn() -> P>,
3025}
3026
3027impl<P, const K: usize, const H: usize> Copy for PatternedOrder2<P, K, H> {}
3028
3029impl<P, const K: usize, const H: usize> Clone for PatternedOrder2<P, K, H> {
3030 fn clone(&self) -> Self {
3031 *self
3032 }
3033}
3034
3035impl<P, const K: usize, const H: usize> PatternedOrder2<P, K, H>
3036where
3037 P: HessianPattern<K, H>,
3038{
3039 #[inline]
3040 #[must_use]
3041 pub fn g(&self) -> [f64; K] {
3042 self.g
3043 }
3044
3045 #[inline]
3049 #[must_use]
3050 pub fn h(&self) -> [[f64; K]; K] {
3051 let mut dense = [[0.0; K]; K];
3052 for (slot, &(i, j)) in P::PAIRS.iter().enumerate() {
3053 dense[i][j] = self.h[slot];
3054 dense[j][i] = self.h[slot];
3055 }
3056 dense
3057 }
3058
3059 #[inline]
3060 fn pair_mask_between(left: u128, right: u128) -> u128 {
3061 let mut result = 0u128;
3062 let mut left_axes = left;
3063 while left_axes != 0 {
3064 let i = left_axes.trailing_zeros() as usize;
3065 left_axes &= left_axes - 1;
3066 let mut right_axes = right;
3067 while right_axes != 0 {
3068 let j = right_axes.trailing_zeros() as usize;
3069 right_axes &= right_axes - 1;
3070 result |= P::PAIR_BITS[i][j];
3071 }
3072 }
3073 result
3074 }
3075}
3076
3077impl<P, const K: usize, const H: usize> JetScalar<K> for PatternedOrder2<P, K, H>
3078where
3079 P: HessianPattern<K, H>,
3080{
3081 #[inline]
3082 fn constant(c: f64) -> Self {
3083 Self {
3084 v: c,
3085 g: [0.0; K],
3086 h: [0.0; H],
3087 gradient_mask: 0,
3088 hessian_mask: 0,
3089 pattern: std::marker::PhantomData,
3090 }
3091 }
3092
3093 #[inline]
3094 fn variable(x: f64, axis: usize) -> Self {
3095 let mut out = Self::constant(x);
3096 if axis < K {
3097 out.g[axis] = 1.0;
3098 out.gradient_mask = 1u128 << axis;
3099 }
3100 out
3101 }
3102}
3103
3104impl<P, const K: usize, const H: usize> crate::nested_dual::JetField for PatternedOrder2<P, K, H>
3105where
3106 P: HessianPattern<K, H>,
3107{
3108 #[inline]
3109 fn value(&self) -> f64 {
3110 self.v
3111 }
3112
3113 #[inline]
3114 fn add(&self, other: &Self) -> Self {
3115 let mut out = Self::constant(self.v + other.v);
3116 out.gradient_mask = self.gradient_mask | other.gradient_mask;
3117 let mut gradient_mask = out.gradient_mask;
3118 while gradient_mask != 0 {
3119 let i = gradient_mask.trailing_zeros() as usize;
3120 gradient_mask &= gradient_mask - 1;
3121 out.g[i] = self.g[i] + other.g[i];
3122 }
3123 out.hessian_mask = self.hessian_mask | other.hessian_mask;
3124 let mut hessian_mask = out.hessian_mask;
3125 while hessian_mask != 0 {
3126 let slot = hessian_mask.trailing_zeros() as usize;
3127 hessian_mask &= hessian_mask - 1;
3128 out.h[slot] = self.h[slot] + other.h[slot];
3129 }
3130 out
3131 }
3132
3133 #[inline]
3134 fn sub(&self, other: &Self) -> Self {
3135 let mut out = Self::constant(self.v - other.v);
3136 out.gradient_mask = self.gradient_mask | other.gradient_mask;
3137 let mut gradient_mask = out.gradient_mask;
3138 while gradient_mask != 0 {
3139 let i = gradient_mask.trailing_zeros() as usize;
3140 gradient_mask &= gradient_mask - 1;
3141 out.g[i] = self.g[i] - other.g[i];
3142 }
3143 out.hessian_mask = self.hessian_mask | other.hessian_mask;
3144 let mut hessian_mask = out.hessian_mask;
3145 while hessian_mask != 0 {
3146 let slot = hessian_mask.trailing_zeros() as usize;
3147 hessian_mask &= hessian_mask - 1;
3148 out.h[slot] = self.h[slot] - other.h[slot];
3149 }
3150 out
3151 }
3152
3153 #[inline]
3154 fn mul(&self, other: &Self) -> Self {
3155 let mut out = Self::constant(self.v * other.v);
3156 out.gradient_mask = self.gradient_mask | other.gradient_mask;
3157 let mut gradient_mask = out.gradient_mask;
3158 while gradient_mask != 0 {
3159 let i = gradient_mask.trailing_zeros() as usize;
3160 gradient_mask &= gradient_mask - 1;
3161 out.g[i] = self.v * other.g[i] + self.g[i] * other.v;
3162 }
3163 out.hessian_mask = self.hessian_mask
3164 | other.hessian_mask
3165 | Self::pair_mask_between(self.gradient_mask, other.gradient_mask);
3166 let mut hessian_mask = out.hessian_mask;
3167 while hessian_mask != 0 {
3168 let slot = hessian_mask.trailing_zeros() as usize;
3169 hessian_mask &= hessian_mask - 1;
3170 let (i, j) = P::PAIRS[slot];
3171 out.h[slot] = self.v * other.h[slot]
3172 + self.g[i] * other.g[j]
3173 + self.g[j] * other.g[i]
3174 + self.h[slot] * other.v;
3175 }
3176 out
3177 }
3178
3179 #[inline]
3180 fn neg(&self) -> Self {
3181 self.scale(-1.0)
3182 }
3183
3184 #[inline]
3185 fn scale(&self, scale: f64) -> Self {
3186 let mut out = Self::constant(self.v * scale);
3187 out.gradient_mask = self.gradient_mask;
3188 let mut gradient_mask = out.gradient_mask;
3189 while gradient_mask != 0 {
3190 let i = gradient_mask.trailing_zeros() as usize;
3191 gradient_mask &= gradient_mask - 1;
3192 out.g[i] = self.g[i] * scale;
3193 }
3194 out.hessian_mask = self.hessian_mask;
3195 let mut hessian_mask = out.hessian_mask;
3196 while hessian_mask != 0 {
3197 let slot = hessian_mask.trailing_zeros() as usize;
3198 hessian_mask &= hessian_mask - 1;
3199 out.h[slot] = self.h[slot] * scale;
3200 }
3201 out
3202 }
3203
3204 #[inline]
3205 fn compose_unary(&self, derivatives: [f64; 5]) -> Self {
3206 let mut out = Self::constant(derivatives[0]);
3207 out.gradient_mask = self.gradient_mask;
3208 let mut gradient_mask = out.gradient_mask;
3209 while gradient_mask != 0 {
3210 let i = gradient_mask.trailing_zeros() as usize;
3211 gradient_mask &= gradient_mask - 1;
3212 out.g[i] = derivatives[1] * self.g[i];
3213 }
3214 out.hessian_mask =
3215 self.hessian_mask | Self::pair_mask_between(self.gradient_mask, self.gradient_mask);
3216 let mut hessian_mask = out.hessian_mask;
3217 while hessian_mask != 0 {
3218 let slot = hessian_mask.trailing_zeros() as usize;
3219 hessian_mask &= hessian_mask - 1;
3220 let (i, j) = P::PAIRS[slot];
3221 out.h[slot] = derivatives[2] * self.g[i] * self.g[j] + derivatives[1] * self.h[slot];
3222 }
3223 out
3224 }
3225}
3226
3227#[derive(Clone, Copy, Debug)]
3240pub struct Order2<const K: usize>(pub crate::jet_tower::Tower2<K>);
3241
3242impl<const K: usize> Order2<K> {
3243 #[inline]
3245 #[must_use]
3246 pub fn g(&self) -> &[f64; K] {
3247 &self.0.g
3248 }
3249
3250 #[inline]
3252 #[must_use]
3253 pub fn h(&self) -> &[[f64; K]; K] {
3254 &self.0.h
3255 }
3256
3257 #[inline]
3259 #[must_use]
3260 pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
3261 let crate::jet_tower::Tower2 { v, g, h } = self.0;
3262 (v, g, h)
3263 }
3264}
3265
3266impl<const K: usize> JetScalar<K> for Order2<K> {
3267 fn constant(c: f64) -> Self {
3268 Order2(crate::jet_tower::Tower2::constant(c))
3269 }
3270 fn variable(x: f64, axis: usize) -> Self {
3271 Order2(crate::jet_tower::Tower2::variable(x, axis))
3272 }
3273
3274 #[inline(always)]
3275 fn symmetric_quadratic_form<C: SymmetricQuadraticCoefficients>(
3276 inputs: &[Self],
3277 coefficients: &C,
3278 ) -> Self {
3279 assert_eq!(inputs.len(), coefficients.dimension());
3280 let input_dimension = inputs.len();
3281 assert!(input_dimension <= K);
3282 let mut values = [0.0; K];
3283 for axis in 0..input_dimension {
3284 values[axis] = inputs[axis].0.v;
3285 }
3286 let mut projected = [0.0; K];
3287 coefficients.multiply(
3288 &values[..input_dimension],
3289 &mut projected[..input_dimension],
3290 );
3291
3292 let mut out = crate::jet_tower::Tower2::zero();
3293 for axis in 0..input_dimension {
3294 out.v += values[axis] * projected[axis];
3295 }
3296 for primary in 0..K {
3297 let mut channel = 0.0;
3298 for axis in 0..input_dimension {
3299 channel += projected[axis] * inputs[axis].0.g[primary];
3300 }
3301 out.g[primary] = 2.0 * channel;
3302 }
3303 let mut input_gradient = [0.0; K];
3304 let mut projected_gradient = [0.0; K];
3305 for primary_b in 0..K {
3306 for row in 0..input_dimension {
3307 input_gradient[row] = inputs[row].0.g[primary_b];
3308 }
3309 coefficients.multiply(
3310 &input_gradient[..input_dimension],
3311 &mut projected_gradient[..input_dimension],
3312 );
3313 for primary_a in 0..=primary_b {
3314 let mut inherited = 0.0;
3315 let mut curvature = 0.0;
3316 for row in 0..input_dimension {
3317 inherited += projected[row] * inputs[row].0.h[primary_a][primary_b];
3318 curvature += inputs[row].0.g[primary_a] * projected_gradient[row];
3319 }
3320 let channel = 2.0 * (inherited + curvature);
3321 out.h[primary_a][primary_b] = channel;
3322 out.h[primary_b][primary_a] = channel;
3323 }
3324 }
3325 Order2(out)
3326 }
3327
3328 #[inline(always)]
3329 fn linear_combination(inputs: &[Self], weights: &[f64]) -> Self {
3330 assert_eq!(inputs.len(), weights.len());
3331 let mut out = crate::jet_tower::Tower2::zero();
3332 for (input, &weight) in inputs.iter().zip(weights) {
3333 out.v += input.0.v * weight;
3334 }
3335 for primary in 0..K {
3336 for (input, &weight) in inputs.iter().zip(weights) {
3337 out.g[primary] += input.0.g[primary] * weight;
3338 }
3339 for other in primary..K {
3340 for (input, &weight) in inputs.iter().zip(weights) {
3341 out.h[primary][other] += input.0.h[primary][other] * weight;
3342 }
3343 out.h[other][primary] = out.h[primary][other];
3344 }
3345 }
3346 Order2(out)
3347 }
3348
3349 #[inline(always)]
3350 fn add_constant(&self, constant: f64) -> Self {
3351 let mut out = *self;
3352 out.0.v += constant;
3353 out
3354 }
3355
3356 #[inline(always)]
3357 fn multiply_add(&self, right: &Self, addend: &Self) -> Self {
3358 let mut out = crate::jet_tower::Tower2::zero();
3359 out.v = self.0.v * right.0.v + addend.0.v;
3360 for primary in 0..K {
3361 out.g[primary] =
3362 self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v + addend.0.g[primary];
3363 for other in primary..K {
3364 let channel = self.0.v * right.0.h[primary][other]
3365 + self.0.g[primary] * right.0.g[other]
3366 + self.0.g[other] * right.0.g[primary]
3367 + self.0.h[primary][other] * right.0.v
3368 + addend.0.h[primary][other];
3369 out.h[primary][other] = channel;
3370 out.h[other][primary] = channel;
3371 }
3372 }
3373 Order2(out)
3374 }
3375
3376 #[inline(always)]
3377 fn product(&self, right: &Self) -> Self {
3378 let mut out = crate::jet_tower::Tower2::zero();
3379 out.v = self.0.v * right.0.v;
3380 for primary in 0..K {
3381 out.g[primary] = self.0.v * right.0.g[primary] + self.0.g[primary] * right.0.v;
3382 for other in primary..K {
3383 let channel = self.0.v * right.0.h[primary][other]
3384 + self.0.g[primary] * right.0.g[other]
3385 + self.0.g[other] * right.0.g[primary]
3386 + self.0.h[primary][other] * right.0.v;
3387 out.h[primary][other] = channel;
3388 out.h[other][primary] = channel;
3389 }
3390 }
3391 Order2(out)
3392 }
3393
3394 #[inline(always)]
3395 fn affine_compose(
3396 &self,
3397 input_scale: f64,
3398 input_shift: f64,
3399 derivative_stack: [f64; 5],
3400 ) -> Self {
3401 assert!(input_shift.is_finite(), "affine input shift must be finite");
3402 let first = derivative_stack[1] * input_scale;
3403 let second = derivative_stack[2] * input_scale * input_scale;
3404 let mut out = crate::jet_tower::Tower2::zero();
3405 out.v = derivative_stack[0];
3406 for primary in 0..K {
3407 out.g[primary] = first * self.0.g[primary];
3408 for other in primary..K {
3409 let channel =
3410 first * self.0.h[primary][other] + second * self.0.g[primary] * self.0.g[other];
3411 out.h[primary][other] = channel;
3412 out.h[other][primary] = channel;
3413 }
3414 }
3415 Order2(out)
3416 }
3417
3418 #[inline(always)]
3419 fn affine_composed_sum(
3420 inputs: &[Self],
3421 input_scales: &[f64],
3422 derivative_stacks: &[[f64; 5]],
3423 ) -> Self {
3424 assert_eq!(inputs.len(), input_scales.len());
3425 assert_eq!(inputs.len(), derivative_stacks.len());
3426 let mut out = crate::jet_tower::Tower2::zero();
3427 for ((input, &input_scale), stack) in inputs.iter().zip(input_scales).zip(derivative_stacks)
3428 {
3429 let first = stack[1] * input_scale;
3430 let second = stack[2] * input_scale * input_scale;
3431 out.v += stack[0];
3432 for primary in 0..K {
3433 out.g[primary] += first * input.0.g[primary];
3434 for other in primary..K {
3435 out.h[primary][other] += first * input.0.h[primary][other]
3436 + second * input.0.g[primary] * input.0.g[other];
3437 }
3438 }
3439 }
3440 for primary in 0..K {
3441 for other in primary + 1..K {
3442 out.h[other][primary] = out.h[primary][other];
3443 }
3444 }
3445 Order2(out)
3446 }
3447
3448 #[inline(always)]
3449 fn shared_multiply_add_affine_composed_sum<const N: usize>(
3450 lefts: &[&Self; N],
3451 right: &Self,
3452 addend: &Self,
3453 addend_scales: &[f64; N],
3454 input_scales: &[f64; N],
3455 derivative_stacks: &[[f64; 5]; N],
3456 ) -> Self {
3457 let (representatives, term_sources, source_count) =
3458 canonical_shared_source_schedule::<N>(|term, representative| {
3459 std::ptr::eq(lefts[term], lefts[representative])
3460 && addend_scales[term] == addend_scales[representative]
3461 });
3462 let (value, source_derivatives) =
3463 aggregate_shared_source_derivatives(&term_sources, input_scales, derivative_stacks);
3464 let mut source_gradients = [[0.0; K]; N];
3465 let mut out = crate::jet_tower::Tower2::zero();
3466 out.v = value;
3467 let mut right_first = 0.0;
3468 let mut addend_first = 0.0;
3469 for source in 0..source_count {
3470 let term = representatives[source];
3471 let first = source_derivatives[source][1];
3472 right_first += first * lefts[term].0.v;
3473 addend_first += first * addend_scales[term];
3474 for primary in 0..K {
3475 let product_gradient =
3476 lefts[term].0.v * right.0.g[primary] + lefts[term].0.g[primary] * right.0.v;
3477 let inner_gradient = if addend_scales[term] == 0.0 {
3478 product_gradient
3479 } else if addend_scales[term] == 1.0 {
3480 product_gradient + addend.0.g[primary]
3481 } else {
3482 product_gradient + addend_scales[term] * addend.0.g[primary]
3483 };
3484 source_gradients[source][primary] = inner_gradient;
3485 out.g[primary] += first * lefts[term].0.g[primary] * right.0.v;
3486 }
3487 }
3488 if N != 0 {
3489 for primary in 0..K {
3490 out.g[primary] += right_first * right.0.g[primary];
3491 }
3492 }
3493 let addend_live = addend_scales.iter().any(|&scale| scale != 0.0);
3494 if addend_live {
3495 for primary in 0..K {
3496 out.g[primary] += addend_first * addend.0.g[primary];
3497 }
3498 }
3499 for primary in 0..K {
3500 for other in primary..K {
3501 let mut channel = if N == 0 {
3502 0.0
3503 } else {
3504 right_first * right.0.h[primary][other]
3505 };
3506 if addend_live {
3507 channel += addend_first * addend.0.h[primary][other];
3508 }
3509 for source in 0..source_count {
3510 let term = representatives[source];
3511 let local_product_hessian = lefts[term].0.g[primary] * right.0.g[other]
3512 + lefts[term].0.g[other] * right.0.g[primary]
3513 + lefts[term].0.h[primary][other] * right.0.v;
3514 channel += source_derivatives[source][1] * local_product_hessian
3515 + source_derivatives[source][2]
3516 * source_gradients[source][primary]
3517 * source_gradients[source][other];
3518 }
3519 out.h[primary][other] = channel;
3520 out.h[other][primary] = channel;
3521 }
3522 }
3523 Order2(out)
3524 }
3525
3526 #[inline(always)]
3527 fn composed_sum(inputs: &[Self], derivative_stacks: &[[f64; 5]]) -> Self {
3528 assert_eq!(inputs.len(), derivative_stacks.len());
3529 let mut out = crate::jet_tower::Tower2::zero();
3530 for (input, stack) in inputs.iter().zip(derivative_stacks) {
3531 out.v += stack[0];
3532 for primary in 0..K {
3533 out.g[primary] += stack[1] * input.0.g[primary];
3534 for other in primary..K {
3535 out.h[primary][other] += stack[1] * input.0.h[primary][other]
3536 + stack[2] * input.0.g[primary] * input.0.g[other];
3537 }
3538 }
3539 }
3540 for primary in 0..K {
3541 for other in primary + 1..K {
3542 out.h[other][primary] = out.h[primary][other];
3543 }
3544 }
3545 Order2(out)
3546 }
3547}
3548
3549impl<const K: usize> crate::nested_dual::JetField for Order2<K> {
3550 fn value(&self) -> f64 {
3551 self.0.v
3552 }
3553 fn add(&self, o: &Self) -> Self {
3554 Order2(self.0 + o.0)
3555 }
3556 fn sub(&self, o: &Self) -> Self {
3557 Order2(self.0 + o.0.scale(-1.0))
3560 }
3561 fn mul(&self, o: &Self) -> Self {
3562 Order2(crate::jet_tower::Tower2::mul(&self.0, &o.0))
3563 }
3564 fn neg(&self) -> Self {
3565 Order2(self.0.scale(-1.0))
3566 }
3567 fn scale(&self, s: f64) -> Self {
3568 Order2(self.0.scale(s))
3569 }
3570 fn compose_unary(&self, d: [f64; 5]) -> Self {
3571 Order2(self.0.compose_unary([d[0], d[1], d[2]]))
3573 }
3574}
3575
3576#[derive(Clone, Copy, Debug)]
3601pub struct MappedOrder2Accumulator<const K: usize> {
3602 value: f64,
3603 gradient: [f64; K],
3604 hessian: [[f64; K]; K],
3605}
3606
3607#[derive(Clone, Copy, Debug)]
3614pub struct StaticOrder2Atom<
3615 const N: usize,
3616 const H: usize,
3617 const GRADIENT_BITS: u128,
3618 const HESSIAN_BITS: u128,
3619> {
3620 value: f64,
3621 gradient: [f64; N],
3622 hessian: [f64; H],
3623}
3624
3625impl<const N: usize, const H: usize, const G: u128, const Q: u128> StaticOrder2Atom<N, H, G, Q> {
3626 #[inline(always)]
3628 #[must_use]
3629 pub fn new(value: f64, gradient: [f64; N], hessian: [f64; H]) -> Self {
3630 assert!(H == N * (N + 1) / 2, "invalid packed order-two shape");
3631 assert!(N <= 128 && H <= 128, "static atom sparsity mask overflow");
3632 Self {
3633 value,
3634 gradient,
3635 hessian,
3636 }
3637 }
3638
3639 #[inline(always)]
3641 #[must_use]
3642 pub fn value(&self) -> f64 {
3643 self.value
3644 }
3645
3646 #[inline(always)]
3648 #[must_use]
3649 pub fn gradient(&self) -> [f64; N] {
3650 self.gradient
3651 }
3652
3653 #[inline(always)]
3655 #[must_use]
3656 pub fn hessian_at(&self, row: usize, column: usize) -> f64 {
3657 assert!(
3658 row < N && column < N,
3659 "static atom Hessian axis out of range"
3660 );
3661 let (row, column) = if row <= column {
3662 (row, column)
3663 } else {
3664 (column, row)
3665 };
3666 let index = row * (2 * N - row + 1) / 2 + column - row;
3667 self.hessian[index]
3668 }
3669}
3670
3671pub trait Order2AtomChannels<const N: usize> {
3676 const GRADIENT_BITS: u128;
3678 const HESSIAN_BITS: u128;
3680 fn gradient_at(&self, axis: usize) -> f64;
3682 fn hessian_at(&self, row: usize, column: usize) -> f64;
3684}
3685
3686impl<const N: usize> Order2AtomChannels<N> for Order2<N> {
3687 const GRADIENT_BITS: u128 = low_mask(N);
3688 const HESSIAN_BITS: u128 = low_mask(N * (N + 1) / 2);
3689
3690 #[inline(always)]
3691 fn gradient_at(&self, axis: usize) -> f64 {
3692 self.0.g[axis]
3693 }
3694
3695 #[inline(always)]
3696 fn hessian_at(&self, row: usize, column: usize) -> f64 {
3697 self.0.h[row][column]
3698 }
3699}
3700
3701impl<const N: usize, const H: usize, const G: u128, const Q: u128> Order2AtomChannels<N>
3702 for StaticOrder2Atom<N, H, G, Q>
3703{
3704 const GRADIENT_BITS: u128 = G;
3705 const HESSIAN_BITS: u128 = Q;
3706
3707 #[inline(always)]
3708 fn gradient_at(&self, axis: usize) -> f64 {
3709 self.gradient[axis]
3710 }
3711
3712 #[inline(always)]
3713 fn hessian_at(&self, row: usize, column: usize) -> f64 {
3714 StaticOrder2Atom::hessian_at(self, row, column)
3715 }
3716}
3717
3718const fn low_mask(channels: usize) -> u128 {
3719 if channels >= 128 {
3720 u128::MAX
3721 } else {
3722 (1u128 << channels) - 1
3723 }
3724}
3725
3726impl<const K: usize> MappedOrder2Accumulator<K> {
3727 #[inline(always)]
3729 #[must_use]
3730 pub fn zero() -> Self {
3731 Self {
3732 value: 0.0,
3733 gradient: [0.0; K],
3734 hessian: [[0.0; K]; K],
3735 }
3736 }
3737
3738 #[inline(always)]
3745 pub fn add_composed<const N: usize, const H: usize, A: Order2AtomChannels<N>>(
3746 &mut self,
3747 atom: &A,
3748 axes: [usize; N],
3749 derivatives: [f64; 3],
3750 value_add: bool,
3751 gradient_add: [bool; N],
3752 hessian_add: [bool; H],
3753 ) {
3754 assert!(H == N * (N + 1) / 2, "invalid mapped Hessian write shape");
3755 assert!(N <= 128 && H <= 128, "mapped atom sparsity mask overflow");
3756 assert!(
3757 axes.iter().all(|&axis| axis < K),
3758 "mapped atom axis must be within the global primary dimension"
3759 );
3760 assert!(
3761 axes.iter()
3762 .enumerate()
3763 .all(|(i, axis)| !axes[..i].contains(axis)),
3764 "mapped atom axes must be injective"
3765 );
3766
3767 if value_add {
3768 self.value += derivatives[0];
3769 } else {
3770 self.value = derivatives[0];
3771 }
3772 let mut packed = 0;
3773 for local_i in 0..N {
3774 let global_i = axes[local_i];
3775 if A::GRADIENT_BITS & (1u128 << local_i) != 0 {
3776 let channel = derivatives[1] * atom.gradient_at(local_i);
3777 if gradient_add[local_i] {
3778 self.gradient[global_i] += channel;
3779 } else {
3780 self.gradient[global_i] = channel;
3781 }
3782 }
3783 for local_j in local_i..N {
3784 let global_j = axes[local_j];
3785 let inner_live = A::HESSIAN_BITS & (1u128 << packed) != 0;
3786 let outer_live = A::GRADIENT_BITS & (1u128 << local_i) != 0
3787 && A::GRADIENT_BITS & (1u128 << local_j) != 0;
3788 let channel = if inner_live {
3789 let inner = derivatives[1] * atom.hessian_at(local_i, local_j);
3790 if outer_live {
3791 inner
3792 + derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
3793 } else {
3794 inner
3795 }
3796 } else if outer_live {
3797 derivatives[2] * atom.gradient_at(local_i) * atom.gradient_at(local_j)
3798 } else {
3799 packed += 1;
3800 continue;
3801 };
3802 if hessian_add[packed] {
3803 self.hessian[global_i][global_j] += channel;
3804 if global_i != global_j {
3805 self.hessian[global_j][global_i] += channel;
3806 }
3807 } else {
3808 self.hessian[global_i][global_j] = channel;
3809 if global_i != global_j {
3810 self.hessian[global_j][global_i] = channel;
3811 }
3812 }
3813 packed += 1;
3814 }
3815 }
3816 }
3817
3818 #[inline(always)]
3820 #[must_use]
3821 pub fn into_channels(self) -> (f64, [f64; K], [[f64; K]; K]) {
3822 (self.value, self.gradient, self.hessian)
3823 }
3824}
3825
3826pub trait DynamicOrder2Term {
3833 fn outer_first(&self) -> f64;
3835
3836 fn outer_second(&self) -> f64;
3838
3839 fn inner_gradient(&self, axis: usize) -> f64;
3841
3842 fn inner_hessian(&self, row: usize, column: usize) -> f64;
3844}
3845
3846#[derive(Debug)]
3863pub struct DynamicOrder2Accumulator {
3864 value: f64,
3865 gradient: Vec<f64>,
3866 hessian: Vec<f64>,
3867}
3868
3869impl DynamicOrder2Accumulator {
3870 #[inline(always)]
3872 #[must_use]
3873 pub fn from_composed_sum<T: DynamicOrder2Term, const N: usize>(
3874 dimension: usize,
3875 value: f64,
3876 terms: &[T; N],
3877 ) -> Self {
3878 let mut gradient = vec![0.0; dimension];
3879 let mut hessian = vec![0.0; dimension * dimension];
3880
3881 for axis in 0..dimension {
3882 let mut channel = 0.0;
3883 for term in terms {
3884 channel += term.outer_first() * term.inner_gradient(axis);
3885 }
3886 gradient[axis] = channel;
3887 }
3888
3889 for row in 0..dimension {
3890 for column in row..dimension {
3891 let mut channel = 0.0;
3892 for term in terms {
3893 let row_gradient = term.inner_gradient(row);
3894 let column_gradient = term.inner_gradient(column);
3895 channel += term.outer_second() * row_gradient * column_gradient
3896 + term.outer_first() * term.inner_hessian(row, column);
3897 }
3898 hessian[row * dimension + column] = channel;
3899 hessian[column * dimension + row] = channel;
3900 }
3901 }
3902
3903 Self {
3904 value,
3905 gradient,
3906 hessian,
3907 }
3908 }
3909
3910 #[inline(always)]
3912 #[must_use]
3913 pub fn into_channels(self) -> (f64, Vec<f64>, Vec<f64>) {
3914 (self.value, self.gradient, self.hessian)
3915 }
3916}
3917
3918pub trait Lane: Copy {
3954 fn splat(x: f64) -> Self;
3956 fn add(self, o: Self) -> Self;
3958 fn sub(self, o: Self) -> Self;
3960 fn mul(self, o: Self) -> Self;
3962 fn lane(self, i: usize) -> f64;
3964 fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3];
3970 fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5];
3979}
3980
3981impl Lane for f64 {
3982 #[inline]
3983 fn splat(x: f64) -> Self {
3984 x
3985 }
3986 #[inline]
3987 fn add(self, o: Self) -> Self {
3988 self + o
3989 }
3990 #[inline]
3991 fn sub(self, o: Self) -> Self {
3992 self - o
3993 }
3994 #[inline]
3995 fn mul(self, o: Self) -> Self {
3996 self * o
3997 }
3998 #[inline]
3999 fn lane(self, _: usize) -> f64 {
4000 self
4001 }
4002 #[inline]
4003 fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
4004 stack(self)
4005 }
4006 #[inline]
4007 fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
4008 stack(self)
4009 }
4010}
4011
4012impl Lane for wide::f64x4 {
4013 #[inline]
4014 fn splat(x: f64) -> Self {
4015 wide::f64x4::splat(x)
4016 }
4017 #[inline]
4018 fn add(self, o: Self) -> Self {
4019 self + o
4020 }
4021 #[inline]
4022 fn sub(self, o: Self) -> Self {
4023 self - o
4024 }
4025 #[inline]
4026 fn mul(self, o: Self) -> Self {
4027 self * o
4028 }
4029 #[inline]
4030 fn lane(self, i: usize) -> f64 {
4031 self.to_array()[i]
4032 }
4033 #[inline]
4034 fn unary3(self, stack: impl Fn(f64) -> [f64; 3]) -> [Self; 3] {
4035 let a = self.to_array();
4036 let mut d0 = [0.0_f64; 4];
4037 let mut d1 = [0.0_f64; 4];
4038 let mut d2 = [0.0_f64; 4];
4039 for i in 0..4 {
4040 let s = stack(a[i]);
4041 d0[i] = s[0];
4042 d1[i] = s[1];
4043 d2[i] = s[2];
4044 }
4045 [
4046 wide::f64x4::new(d0),
4047 wide::f64x4::new(d1),
4048 wide::f64x4::new(d2),
4049 ]
4050 }
4051 #[inline]
4052 fn unary5(self, stack: impl Fn(f64) -> [f64; 5]) -> [Self; 5] {
4053 let a = self.to_array();
4054 let mut d = [[0.0_f64; 4]; 5];
4055 for i in 0..4 {
4056 let s = stack(a[i]);
4057 for (k, dk) in d.iter_mut().enumerate() {
4058 dk[i] = s[k];
4059 }
4060 }
4061 [
4062 wide::f64x4::new(d[0]),
4063 wide::f64x4::new(d[1]),
4064 wide::f64x4::new(d[2]),
4065 wide::f64x4::new(d[3]),
4066 wide::f64x4::new(d[4]),
4067 ]
4068 }
4069}
4070
4071#[derive(Clone, Copy, Debug)]
4081pub struct Order2Lane<L: Lane, const K: usize> {
4082 pub v: L,
4084 pub g: [L; K],
4086 pub h: [[L; K]; K],
4088}
4089
4090pub type Order2Batch<const K: usize> = Order2Lane<wide::f64x4, K>;
4092
4093impl<L: Lane, const K: usize> Order2Lane<L, K> {
4094 #[inline]
4096 pub fn constant(c: L) -> Self {
4097 Order2Lane {
4098 v: c,
4099 g: [L::splat(0.0); K],
4100 h: [[L::splat(0.0); K]; K],
4101 }
4102 }
4103
4104 #[inline]
4108 pub fn variable(value: L, axis: usize) -> Self {
4109 let mut out = Self::constant(value);
4110 out.g[axis] = L::splat(1.0);
4111 out
4112 }
4113
4114 #[inline]
4116 pub fn add(&self, o: &Self) -> Self {
4117 let mut out = *self;
4118 out.v = self.v.add(o.v);
4119 for i in 0..K {
4120 out.g[i] = self.g[i].add(o.g[i]);
4121 for j in 0..K {
4122 out.h[i][j] = self.h[i][j].add(o.h[i][j]);
4123 }
4124 }
4125 out
4126 }
4127
4128 #[inline]
4130 pub fn scale(&self, s: f64) -> Self {
4131 let sl = L::splat(s);
4132 let mut out = *self;
4133 out.v = self.v.mul(sl);
4134 for i in 0..K {
4135 out.g[i] = self.g[i].mul(sl);
4136 for j in 0..K {
4137 out.h[i][j] = self.h[i][j].mul(sl);
4138 }
4139 }
4140 out
4141 }
4142
4143 #[inline]
4146 pub fn sub(&self, o: &Self) -> Self {
4147 self.add(&o.scale(-1.0))
4148 }
4149
4150 #[inline]
4152 pub fn neg(&self) -> Self {
4153 self.scale(-1.0)
4154 }
4155
4156 #[inline]
4171 pub fn mul(&self, o: &Self) -> Self {
4172 let a = self;
4173 let b = o;
4174 let mut out = Self::constant(a.v.mul(b.v));
4175 for i in 0..K {
4176 out.g[i] = a.v.mul(b.g[i]).add(a.g[i].mul(b.v));
4178 }
4179 for i in 0..K {
4180 for j in i..K {
4181 let hij =
4183 a.v.mul(b.h[i][j])
4184 .add(a.g[i].mul(b.g[j]))
4185 .add(a.g[j].mul(b.g[i]))
4186 .add(a.h[i][j].mul(b.v));
4187 out.h[i][j] = hij;
4188 out.h[j][i] = hij;
4189 }
4190 }
4191 out
4192 }
4193
4194 #[inline]
4199 pub fn compose_unary(&self, d: [L; 3]) -> Self {
4200 let mut out = Self::constant(d[0]);
4201 for i in 0..K {
4202 let mut acc = L::splat(0.0);
4203 acc = acc.add(d[1].mul(self.g[i]));
4204 out.g[i] = acc;
4205 }
4206 for i in 0..K {
4207 for j in 0..K {
4208 let mut acc = L::splat(0.0);
4209 acc = acc.add(d[1].mul(self.h[i][j]));
4210 acc = acc.add(d[2].mul(self.g[i]).mul(self.g[j]));
4211 out.h[i][j] = acc;
4212 }
4213 }
4214 out
4215 }
4216
4217 #[inline]
4220 pub fn exp(&self) -> Self {
4221 let d = self.v.unary3(|u| {
4222 let e = u.exp();
4223 [e, e, e]
4224 });
4225 self.compose_unary(d)
4226 }
4227
4228 #[inline]
4231 pub fn ln(&self) -> Self {
4232 let d = self.v.unary3(|u| {
4233 let r = 1.0 / u;
4234 [u.ln(), r, -r * r]
4235 });
4236 self.compose_unary(d)
4237 }
4238
4239 #[inline]
4242 pub fn sqrt(&self) -> Self {
4243 let d = self.v.unary3(|u| {
4244 let s = u.sqrt();
4245 [s, 0.5 / s, -0.25 / (u * s)]
4246 });
4247 self.compose_unary(d)
4248 }
4249
4250 #[inline]
4252 pub fn recip(&self) -> Self {
4253 let d = self.v.unary3(|u| {
4254 let r = 1.0 / u;
4255 let r2 = r * r;
4256 [r, -r2, 2.0 * r2 * r]
4257 });
4258 self.compose_unary(d)
4259 }
4260
4261 #[inline]
4264 pub fn powf(&self, a: f64) -> Self {
4265 let d = self.v.unary3(|u| {
4266 [
4267 u.powf(a),
4268 a * u.powf(a - 1.0),
4269 a * (a - 1.0) * u.powf(a - 2.0),
4270 ]
4271 });
4272 self.compose_unary(d)
4273 }
4274}
4275
4276impl<const K: usize> Order2Batch<K> {
4277 #[inline]
4281 #[must_use]
4282 pub fn lane(&self, i: usize) -> Order2<K> {
4283 let mut t = crate::jet_tower::Tower2::<K>::constant(self.v.lane(i));
4284 for a in 0..K {
4285 t.g[a] = self.g[a].lane(i);
4286 for b in 0..K {
4287 t.h[a][b] = self.h[a][b].lane(i);
4288 }
4289 }
4290 Order2(t)
4291 }
4292}
4293
4294#[derive(Clone, Copy, Debug)]
4310pub struct Order1<const K: usize> {
4311 pub v: f64,
4313 pub g: [f64; K],
4315}
4316
4317impl<const K: usize> Order1<K> {
4318 #[inline]
4320 #[must_use]
4321 pub fn g(&self) -> &[f64; K] {
4322 &self.g
4323 }
4324
4325 #[inline]
4327 #[must_use]
4328 pub fn into_channels(self) -> (f64, [f64; K]) {
4329 (self.v, self.g)
4330 }
4331}
4332
4333impl<const K: usize> JetScalar<K> for Order1<K> {
4334 fn constant(c: f64) -> Self {
4335 Order1 { v: c, g: [0.0; K] }
4337 }
4338 fn variable(x: f64, axis: usize) -> Self {
4339 let mut g = [0.0; K];
4341 g[axis] = 1.0;
4342 Order1 { v: x, g }
4343 }
4344}
4345
4346impl<const K: usize> crate::nested_dual::JetField for Order1<K> {
4347 fn value(&self) -> f64 {
4348 self.v
4349 }
4350 fn add(&self, o: &Self) -> Self {
4351 let mut g = self.g;
4353 for i in 0..K {
4354 g[i] += o.g[i];
4355 }
4356 Order1 { v: self.v + o.v, g }
4357 }
4358 fn sub(&self, o: &Self) -> Self {
4359 self.add(&o.scale(-1.0))
4361 }
4362 fn mul(&self, o: &Self) -> Self {
4363 let a = self;
4368 let b = o;
4369 let mut g = [0.0; K];
4370 for i in 0..K {
4371 g[i] = a.v * b.g[i] + a.g[i] * b.v;
4372 }
4373 Order1 { v: a.v * b.v, g }
4374 }
4375 fn neg(&self) -> Self {
4376 self.scale(-1.0)
4378 }
4379 fn scale(&self, s: f64) -> Self {
4380 let mut g = self.g;
4382 for i in 0..K {
4383 g[i] *= s;
4384 }
4385 Order1 { v: self.v * s, g }
4386 }
4387 fn compose_unary(&self, d: [f64; 5]) -> Self {
4388 let mut g = [0.0; K];
4394 for i in 0..K {
4395 g[i] = d[1] * self.g[i];
4396 }
4397 Order1 { v: d[0], g }
4398 }
4399}
4400
4401#[derive(Clone, Copy, Debug)]
4418pub struct OneSeed<const K: usize> {
4419 pub base: Order2<K>,
4421 pub eps: Order2<K>,
4424}
4425
4426impl<const K: usize> OneSeed<K> {
4427 pub fn seed_direction(x: f64, axis: usize, u_axis: f64) -> Self {
4431 OneSeed {
4432 base: Order2::variable(x, axis),
4433 eps: Order2::constant(u_axis),
4434 }
4435 }
4436
4437 pub fn contracted_third(&self) -> [[f64; K]; K] {
4440 *self.eps.h()
4441 }
4442}
4443
4444impl<const K: usize> JetScalar<K> for OneSeed<K> {
4445 fn constant(c: f64) -> Self {
4446 OneSeed {
4447 base: Order2::constant(c),
4448 eps: Order2::constant(0.0),
4449 }
4450 }
4451 fn variable(x: f64, axis: usize) -> Self {
4452 OneSeed {
4454 base: Order2::variable(x, axis),
4455 eps: Order2::constant(0.0),
4456 }
4457 }
4458}
4459
4460impl<const K: usize> crate::nested_dual::JetField for OneSeed<K> {
4461 fn value(&self) -> f64 {
4462 self.base.value()
4463 }
4464 fn add(&self, o: &Self) -> Self {
4465 OneSeed {
4466 base: self.base.add(&o.base),
4467 eps: self.eps.add(&o.eps),
4468 }
4469 }
4470 fn sub(&self, o: &Self) -> Self {
4471 OneSeed {
4472 base: self.base.sub(&o.base),
4473 eps: self.eps.sub(&o.eps),
4474 }
4475 }
4476 fn mul(&self, o: &Self) -> Self {
4477 let ab = &self.base.0;
4483 let ae = &self.eps.0;
4484 let bb = &o.base.0;
4485 let be = &o.eps.0;
4486 let mut eps = crate::jet_tower::Tower2::<K>::zero();
4487 eps.v = ab.v * be.v + ae.v * bb.v;
4488 for i in 0..K {
4489 eps.g[i] = ab.v * be.g[i] + ab.g[i] * be.v + ae.v * bb.g[i] + ae.g[i] * bb.v;
4490 }
4491 for i in 0..K {
4492 for j in i..K {
4493 let channel = ab.v * be.h[i][j]
4494 + ab.g[i] * be.g[j]
4495 + ab.g[j] * be.g[i]
4496 + ab.h[i][j] * be.v
4497 + ae.v * bb.h[i][j]
4498 + ae.g[i] * bb.g[j]
4499 + ae.g[j] * bb.g[i]
4500 + ae.h[i][j] * bb.v;
4501 eps.h[i][j] = channel;
4502 eps.h[j][i] = channel;
4503 }
4504 }
4505 OneSeed {
4506 base: self.base.mul(&o.base),
4507 eps: Order2(eps),
4508 }
4509 }
4510 fn neg(&self) -> Self {
4511 OneSeed {
4512 base: self.base.neg(),
4513 eps: self.eps.neg(),
4514 }
4515 }
4516 fn scale(&self, s: f64) -> Self {
4517 OneSeed {
4518 base: self.base.scale(s),
4519 eps: self.eps.scale(s),
4520 }
4521 }
4522 fn compose_unary(&self, d: [f64; 5]) -> Self {
4523 let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
4528 let b = &self.base.0;
4529 let e = &self.eps.0;
4530 let mut eps = crate::jet_tower::Tower2::<K>::zero();
4531 eps.v = d[1] * e.v;
4532 for i in 0..K {
4533 eps.g[i] = d[2] * b.g[i] * e.v + d[1] * e.g[i];
4534 }
4535 for i in 0..K {
4536 for j in i..K {
4537 let channel = d[1] * e.h[i][j]
4538 + d[2] * (b.g[i] * e.g[j] + b.g[j] * e.g[i] + b.h[i][j] * e.v)
4539 + d[3] * b.g[i] * b.g[j] * e.v;
4540 eps.h[i][j] = channel;
4541 eps.h[j][i] = channel;
4542 }
4543 }
4544 OneSeed {
4545 base,
4546 eps: Order2(eps),
4547 }
4548 }
4549}
4550
4551#[derive(Clone, Copy, Debug)]
4563pub struct OneSeedLane<L: Lane, const K: usize> {
4564 pub base: Order2Lane<L, K>,
4566 pub eps: Order2Lane<L, K>,
4569}
4570
4571pub type OneSeedBatch<const K: usize> = OneSeedLane<wide::f64x4, K>;
4573
4574impl<L: Lane, const K: usize> OneSeedLane<L, K> {
4575 #[inline]
4577 pub fn constant(c: L) -> Self {
4578 OneSeedLane {
4579 base: Order2Lane::constant(c),
4580 eps: Order2Lane::constant(L::splat(0.0)),
4581 }
4582 }
4583
4584 #[inline]
4587 pub fn variable(value: L, axis: usize) -> Self {
4588 OneSeedLane {
4589 base: Order2Lane::variable(value, axis),
4590 eps: Order2Lane::constant(L::splat(0.0)),
4591 }
4592 }
4593
4594 #[inline]
4599 pub fn seed_direction(value: L, axis: usize, u_axis: L) -> Self {
4600 OneSeedLane {
4601 base: Order2Lane::variable(value, axis),
4602 eps: Order2Lane::constant(u_axis),
4603 }
4604 }
4605
4606 #[inline]
4609 #[must_use]
4610 pub fn contracted_third(&self) -> [[L; K]; K] {
4611 self.eps.h
4612 }
4613
4614 #[inline]
4616 pub fn add(&self, o: &Self) -> Self {
4617 OneSeedLane {
4618 base: self.base.add(&o.base),
4619 eps: self.eps.add(&o.eps),
4620 }
4621 }
4622
4623 #[inline]
4625 pub fn sub(&self, o: &Self) -> Self {
4626 OneSeedLane {
4627 base: self.base.sub(&o.base),
4628 eps: self.eps.sub(&o.eps),
4629 }
4630 }
4631
4632 #[inline]
4634 pub fn mul(&self, o: &Self) -> Self {
4635 let ab = &self.base;
4636 let ae = &self.eps;
4637 let bb = &o.base;
4638 let be = &o.eps;
4639 let mut eps = Order2Lane::constant(ab.v.mul(be.v).add(ae.v.mul(bb.v)));
4640 for i in 0..K {
4641 eps.g[i] =
4642 ab.v.mul(be.g[i])
4643 .add(ab.g[i].mul(be.v))
4644 .add(ae.v.mul(bb.g[i]))
4645 .add(ae.g[i].mul(bb.v));
4646 }
4647 for i in 0..K {
4648 for j in i..K {
4649 let channel =
4650 ab.v.mul(be.h[i][j])
4651 .add(ab.g[i].mul(be.g[j]))
4652 .add(ab.g[j].mul(be.g[i]))
4653 .add(ab.h[i][j].mul(be.v))
4654 .add(ae.v.mul(bb.h[i][j]))
4655 .add(ae.g[i].mul(bb.g[j]))
4656 .add(ae.g[j].mul(bb.g[i]))
4657 .add(ae.h[i][j].mul(bb.v));
4658 eps.h[i][j] = channel;
4659 eps.h[j][i] = channel;
4660 }
4661 }
4662 OneSeedLane {
4663 base: self.base.mul(&o.base),
4664 eps,
4665 }
4666 }
4667
4668 #[inline]
4670 pub fn neg(&self) -> Self {
4671 OneSeedLane {
4672 base: self.base.neg(),
4673 eps: self.eps.neg(),
4674 }
4675 }
4676
4677 #[inline]
4679 pub fn scale(&self, s: f64) -> Self {
4680 OneSeedLane {
4681 base: self.base.scale(s),
4682 eps: self.eps.scale(s),
4683 }
4684 }
4685
4686 #[inline]
4692 pub fn compose_unary(&self, d: [L; 5]) -> Self {
4693 let base = self.base.compose_unary([d[0], d[1], d[2]]);
4694 let b = &self.base;
4695 let e = &self.eps;
4696 let mut eps = Order2Lane::constant(d[1].mul(e.v));
4697 for i in 0..K {
4698 eps.g[i] = d[2].mul(b.g[i]).mul(e.v).add(d[1].mul(e.g[i]));
4699 }
4700 for i in 0..K {
4701 for j in i..K {
4702 let mixed = b.g[i]
4703 .mul(e.g[j])
4704 .add(b.g[j].mul(e.g[i]))
4705 .add(b.h[i][j].mul(e.v));
4706 let channel = d[1]
4707 .mul(e.h[i][j])
4708 .add(d[2].mul(mixed))
4709 .add(d[3].mul(b.g[i]).mul(b.g[j]).mul(e.v));
4710 eps.h[i][j] = channel;
4711 eps.h[j][i] = channel;
4712 }
4713 }
4714 OneSeedLane { base, eps }
4715 }
4716
4717 #[inline]
4719 pub fn exp(&self) -> Self {
4720 let d = self.base.v.unary5(|u| {
4721 let e = u.exp();
4722 [e, e, e, e, e]
4723 });
4724 self.compose_unary(d)
4725 }
4726
4727 #[inline]
4729 pub fn ln(&self) -> Self {
4730 let d = self.base.v.unary5(|u| {
4731 let r = 1.0 / u;
4732 [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
4733 });
4734 self.compose_unary(d)
4735 }
4736
4737 #[inline]
4739 pub fn sqrt(&self) -> Self {
4740 let d = self.base.v.unary5(|u| {
4741 let s = u.sqrt();
4742 [
4743 s,
4744 0.5 / s,
4745 -0.25 / (u * s),
4746 0.375 / (u * u * s),
4747 -0.9375 / (u * u * u * s),
4748 ]
4749 });
4750 self.compose_unary(d)
4751 }
4752
4753 #[inline]
4755 pub fn recip(&self) -> Self {
4756 let d = self.base.v.unary5(|u| {
4757 let r = 1.0 / u;
4758 let r2 = r * r;
4759 [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
4760 });
4761 self.compose_unary(d)
4762 }
4763
4764 #[inline]
4767 pub fn powf(&self, a: f64) -> Self {
4768 let d = self.base.v.unary5(|u| {
4769 [
4770 u.powf(a),
4771 a * u.powf(a - 1.0),
4772 a * (a - 1.0) * u.powf(a - 2.0),
4773 a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
4774 a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
4775 ]
4776 });
4777 self.compose_unary(d)
4778 }
4779
4780 #[inline]
4783 pub fn ln_gamma(&self) -> Self {
4784 let d = self
4785 .base
4786 .v
4787 .unary5(crate::jet_tower::ln_gamma_derivative_stack);
4788 self.compose_unary(d)
4789 }
4790
4791 #[inline]
4794 pub fn digamma(&self) -> Self {
4795 let d = self
4796 .base
4797 .v
4798 .unary5(crate::jet_tower::digamma_derivative_stack);
4799 self.compose_unary(d)
4800 }
4801}
4802
4803impl<const K: usize> OneSeedBatch<K> {
4804 #[inline]
4808 #[must_use]
4809 pub fn lane(&self, i: usize) -> OneSeed<K> {
4810 OneSeed {
4811 base: self.base.lane(i),
4812 eps: self.eps.lane(i),
4813 }
4814 }
4815}
4816
4817#[derive(Clone, Copy, Debug)]
4834pub struct TwoSeed<const K: usize> {
4835 pub base: Order2<K>,
4837 pub eps: Order2<K>,
4839 pub del: Order2<K>,
4841 pub eps_del: Order2<K>,
4844}
4845
4846impl<const K: usize> TwoSeed<K> {
4847 pub fn seed(x: f64, axis: usize, u_axis: f64, v_axis: f64) -> Self {
4851 TwoSeed {
4852 base: Order2::variable(x, axis),
4853 eps: Order2::constant(u_axis),
4854 del: Order2::constant(v_axis),
4855 eps_del: Order2::constant(0.0),
4856 }
4857 }
4858
4859 pub fn contracted_fourth(&self) -> [[f64; K]; K] {
4862 *self.eps_del.h()
4863 }
4864}
4865
4866impl<const K: usize> JetScalar<K> for TwoSeed<K> {
4867 fn constant(c: f64) -> Self {
4868 TwoSeed {
4869 base: Order2::constant(c),
4870 eps: Order2::constant(0.0),
4871 del: Order2::constant(0.0),
4872 eps_del: Order2::constant(0.0),
4873 }
4874 }
4875 fn variable(x: f64, axis: usize) -> Self {
4876 TwoSeed {
4877 base: Order2::variable(x, axis),
4878 eps: Order2::constant(0.0),
4879 del: Order2::constant(0.0),
4880 eps_del: Order2::constant(0.0),
4881 }
4882 }
4883}
4884
4885impl<const K: usize> crate::nested_dual::JetField for TwoSeed<K> {
4886 fn value(&self) -> f64 {
4887 self.base.value()
4888 }
4889 fn add(&self, o: &Self) -> Self {
4890 TwoSeed {
4891 base: self.base.add(&o.base),
4892 eps: self.eps.add(&o.eps),
4893 del: self.del.add(&o.del),
4894 eps_del: self.eps_del.add(&o.eps_del),
4895 }
4896 }
4897 fn sub(&self, o: &Self) -> Self {
4898 TwoSeed {
4899 base: self.base.sub(&o.base),
4900 eps: self.eps.sub(&o.eps),
4901 del: self.del.sub(&o.del),
4902 eps_del: self.eps_del.sub(&o.eps_del),
4903 }
4904 }
4905 fn mul(&self, o: &Self) -> Self {
4906 let a = self;
4907 let b = o;
4908 let base = a.base.mul(&b.base);
4910 let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
4911 let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
4912 let eps_del = a
4913 .base
4914 .mul(&b.eps_del)
4915 .add(&a.eps.mul(&b.del))
4916 .add(&a.del.mul(&b.eps))
4917 .add(&a.eps_del.mul(&b.base));
4918 TwoSeed {
4919 base,
4920 eps,
4921 del,
4922 eps_del,
4923 }
4924 }
4925 fn neg(&self) -> Self {
4926 TwoSeed {
4927 base: self.base.neg(),
4928 eps: self.eps.neg(),
4929 del: self.del.neg(),
4930 eps_del: self.eps_del.neg(),
4931 }
4932 }
4933 fn scale(&self, s: f64) -> Self {
4934 TwoSeed {
4935 base: self.base.scale(s),
4936 eps: self.eps.scale(s),
4937 del: self.del.scale(s),
4938 eps_del: self.eps_del.scale(s),
4939 }
4940 }
4941 fn compose_unary(&self, d: [f64; 5]) -> Self {
4942 let base = self.base.compose_unary([d[0], d[1], d[2], d[3], d[4]]);
4952 let fprime = self.base.compose_unary([d[1], d[2], d[3], d[4], d[4]]); let fsecond = self.base.compose_unary([d[2], d[3], d[4], d[4], d[4]]); let eps = fprime.mul(&self.eps);
4955 let del = fprime.mul(&self.del);
4956 let eps_del = fsecond
4957 .mul(&self.eps)
4958 .mul(&self.del)
4959 .add(&fprime.mul(&self.eps_del));
4960 TwoSeed {
4961 base,
4962 eps,
4963 del,
4964 eps_del,
4965 }
4966 }
4967}
4968
4969#[derive(Clone, Copy, Debug)]
4980pub struct TwoSeedLane<L: Lane, const K: usize> {
4981 pub base: Order2Lane<L, K>,
4983 pub eps: Order2Lane<L, K>,
4985 pub del: Order2Lane<L, K>,
4987 pub eps_del: Order2Lane<L, K>,
4990}
4991
4992pub type TwoSeedBatch<const K: usize> = TwoSeedLane<wide::f64x4, K>;
4994
4995impl<L: Lane, const K: usize> TwoSeedLane<L, K> {
4996 #[inline]
4999 pub fn constant(c: L) -> Self {
5000 let z = Order2Lane::constant(L::splat(0.0));
5001 TwoSeedLane {
5002 base: Order2Lane::constant(c),
5003 eps: z,
5004 del: z,
5005 eps_del: z,
5006 }
5007 }
5008
5009 #[inline]
5012 pub fn variable(value: L, axis: usize) -> Self {
5013 let z = Order2Lane::constant(L::splat(0.0));
5014 TwoSeedLane {
5015 base: Order2Lane::variable(value, axis),
5016 eps: z,
5017 del: z,
5018 eps_del: z,
5019 }
5020 }
5021
5022 #[inline]
5026 pub fn seed(value: L, axis: usize, u_axis: L, v_axis: L) -> Self {
5027 TwoSeedLane {
5028 base: Order2Lane::variable(value, axis),
5029 eps: Order2Lane::constant(u_axis),
5030 del: Order2Lane::constant(v_axis),
5031 eps_del: Order2Lane::constant(L::splat(0.0)),
5032 }
5033 }
5034
5035 #[inline]
5039 #[must_use]
5040 pub fn contracted_fourth(&self) -> [[L; K]; K] {
5041 self.eps_del.h
5042 }
5043
5044 #[inline]
5046 pub fn add(&self, o: &Self) -> Self {
5047 TwoSeedLane {
5048 base: self.base.add(&o.base),
5049 eps: self.eps.add(&o.eps),
5050 del: self.del.add(&o.del),
5051 eps_del: self.eps_del.add(&o.eps_del),
5052 }
5053 }
5054
5055 #[inline]
5057 pub fn sub(&self, o: &Self) -> Self {
5058 TwoSeedLane {
5059 base: self.base.sub(&o.base),
5060 eps: self.eps.sub(&o.eps),
5061 del: self.del.sub(&o.del),
5062 eps_del: self.eps_del.sub(&o.eps_del),
5063 }
5064 }
5065
5066 #[inline]
5068 pub fn mul(&self, o: &Self) -> Self {
5069 let a = self;
5070 let b = o;
5071 let base = a.base.mul(&b.base);
5072 let eps = a.base.mul(&b.eps).add(&a.eps.mul(&b.base));
5073 let del = a.base.mul(&b.del).add(&a.del.mul(&b.base));
5074 let eps_del = a
5075 .base
5076 .mul(&b.eps_del)
5077 .add(&a.eps.mul(&b.del))
5078 .add(&a.del.mul(&b.eps))
5079 .add(&a.eps_del.mul(&b.base));
5080 TwoSeedLane {
5081 base,
5082 eps,
5083 del,
5084 eps_del,
5085 }
5086 }
5087
5088 #[inline]
5090 pub fn neg(&self) -> Self {
5091 TwoSeedLane {
5092 base: self.base.neg(),
5093 eps: self.eps.neg(),
5094 del: self.del.neg(),
5095 eps_del: self.eps_del.neg(),
5096 }
5097 }
5098
5099 #[inline]
5101 pub fn scale(&self, s: f64) -> Self {
5102 TwoSeedLane {
5103 base: self.base.scale(s),
5104 eps: self.eps.scale(s),
5105 del: self.del.scale(s),
5106 eps_del: self.eps_del.scale(s),
5107 }
5108 }
5109
5110 #[inline]
5116 pub fn compose_unary(&self, d: [L; 5]) -> Self {
5117 let base = self.base.compose_unary([d[0], d[1], d[2]]);
5118 let fprime = self.base.compose_unary([d[1], d[2], d[3]]);
5119 let fsecond = self.base.compose_unary([d[2], d[3], d[4]]);
5120 let eps = fprime.mul(&self.eps);
5121 let del = fprime.mul(&self.del);
5122 let eps_del = fsecond
5123 .mul(&self.eps)
5124 .mul(&self.del)
5125 .add(&fprime.mul(&self.eps_del));
5126 TwoSeedLane {
5127 base,
5128 eps,
5129 del,
5130 eps_del,
5131 }
5132 }
5133
5134 #[inline]
5136 pub fn exp(&self) -> Self {
5137 let d = self.base.v.unary5(|u| {
5138 let e = u.exp();
5139 [e, e, e, e, e]
5140 });
5141 self.compose_unary(d)
5142 }
5143
5144 #[inline]
5146 pub fn ln(&self) -> Self {
5147 let d = self.base.v.unary5(|u| {
5148 let r = 1.0 / u;
5149 [u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r]
5150 });
5151 self.compose_unary(d)
5152 }
5153
5154 #[inline]
5156 pub fn sqrt(&self) -> Self {
5157 let d = self.base.v.unary5(|u| {
5158 let s = u.sqrt();
5159 [
5160 s,
5161 0.5 / s,
5162 -0.25 / (u * s),
5163 0.375 / (u * u * s),
5164 -0.9375 / (u * u * u * s),
5165 ]
5166 });
5167 self.compose_unary(d)
5168 }
5169
5170 #[inline]
5172 pub fn recip(&self) -> Self {
5173 let d = self.base.v.unary5(|u| {
5174 let r = 1.0 / u;
5175 let r2 = r * r;
5176 [r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r]
5177 });
5178 self.compose_unary(d)
5179 }
5180
5181 #[inline]
5184 pub fn powf(&self, a: f64) -> Self {
5185 let d = self.base.v.unary5(|u| {
5186 [
5187 u.powf(a),
5188 a * u.powf(a - 1.0),
5189 a * (a - 1.0) * u.powf(a - 2.0),
5190 a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0),
5191 a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0),
5192 ]
5193 });
5194 self.compose_unary(d)
5195 }
5196
5197 #[inline]
5199 pub fn ln_gamma(&self) -> Self {
5200 let d = self
5201 .base
5202 .v
5203 .unary5(crate::jet_tower::ln_gamma_derivative_stack);
5204 self.compose_unary(d)
5205 }
5206
5207 #[inline]
5210 pub fn digamma(&self) -> Self {
5211 let d = self
5212 .base
5213 .v
5214 .unary5(crate::jet_tower::digamma_derivative_stack);
5215 self.compose_unary(d)
5216 }
5217}
5218
5219impl<const K: usize> TwoSeedBatch<K> {
5220 #[inline]
5224 #[must_use]
5225 pub fn lane(&self, i: usize) -> TwoSeed<K> {
5226 TwoSeed {
5227 base: self.base.lane(i),
5228 eps: self.eps.lane(i),
5229 del: self.del.lane(i),
5230 eps_del: self.eps_del.lane(i),
5231 }
5232 }
5233}
5234
5235impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower3<K> {
5242 fn constant(c: f64) -> Self {
5243 crate::jet_tower::Tower3::constant(c)
5244 }
5245 fn variable(x: f64, axis: usize) -> Self {
5246 crate::jet_tower::Tower3::variable(x, axis)
5247 }
5248}
5249
5250impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower3<K> {
5251 fn value(&self) -> f64 {
5252 self.v
5253 }
5254 fn add(&self, o: &Self) -> Self {
5255 *self + *o
5256 }
5257 fn sub(&self, o: &Self) -> Self {
5258 *self + o.scale(-1.0)
5259 }
5260 fn mul(&self, o: &Self) -> Self {
5261 crate::jet_tower::Tower3::mul(self, o)
5262 }
5263 fn neg(&self) -> Self {
5264 self.scale(-1.0)
5265 }
5266 fn scale(&self, s: f64) -> Self {
5267 crate::jet_tower::Tower3::scale(self, s)
5268 }
5269 fn compose_unary(&self, d: [f64; 5]) -> Self {
5270 crate::jet_tower::Tower3::compose_unary(self, [d[0], d[1], d[2], d[3]])
5271 }
5272}
5273
5274impl<const K: usize> JetScalar<K> for crate::jet_tower::Tower4<K> {
5289 fn constant(c: f64) -> Self {
5290 crate::jet_tower::Tower4::constant(c)
5291 }
5292 fn variable(x: f64, axis: usize) -> Self {
5293 crate::jet_tower::Tower4::variable(x, axis)
5294 }
5295}
5296
5297impl<const K: usize> crate::nested_dual::JetField for crate::jet_tower::Tower4<K> {
5298 fn value(&self) -> f64 {
5299 self.v
5300 }
5301 fn add(&self, o: &Self) -> Self {
5302 *self + *o
5303 }
5304 fn sub(&self, o: &Self) -> Self {
5305 *self - *o
5306 }
5307 fn mul(&self, o: &Self) -> Self {
5308 crate::jet_tower::Tower4::mul(self, o)
5309 }
5310 fn neg(&self) -> Self {
5311 self.scale(-1.0)
5312 }
5313 fn scale(&self, s: f64) -> Self {
5314 crate::jet_tower::Tower4::scale(self, s)
5315 }
5316 fn compose_unary(&self, d: [f64; 5]) -> Self {
5317 crate::jet_tower::Tower4::compose_unary(self, d)
5318 }
5319}
5320
5321#[cfg(test)]
5322mod tests {
5323 use super::*;
5324 use crate::jet_tower::{RowProgram, Tower4, program_full_tower};
5325 use crate::nested_dual::JetField;
5326
5327 struct DenseSymmetric3([[f64; 3]; 3]);
5328
5329 impl SymmetricQuadraticCoefficients for DenseSymmetric3 {
5330 fn dimension(&self) -> usize {
5331 3
5332 }
5333
5334 fn multiply(&self, input: &[f64], output: &mut [f64]) {
5335 assert_eq!(input.len(), 3);
5336 assert_eq!(output.len(), 3);
5337 for (row, output) in output.iter_mut().enumerate() {
5338 *output = (0..3)
5339 .map(|column| self.0[row][column] * input[column])
5340 .sum();
5341 }
5342 }
5343
5344 fn coefficient(&self, row: usize, column: usize) -> f64 {
5345 self.0[row][column]
5346 }
5347 }
5348
5349 #[test]
5350 fn symmetric_quadratic_order2_lowerings_match_scalar_program() {
5351 const K: usize = 4;
5352 let coefficients = DenseSymmetric3([[1.2, 0.3, -0.2], [0.3, 0.8, 0.15], [-0.2, 0.15, 1.5]]);
5353 let values = [0.4, -0.7, 1.1, 0.25];
5354 let fixed_vars: [Order2<K>; K] =
5355 std::array::from_fn(|axis| Order2::variable(values[axis], axis));
5356 let fixed_inputs = [
5357 fixed_vars[0].mul(&fixed_vars[1]).add(&fixed_vars[3]),
5358 fixed_vars[1].exp().add(&fixed_vars[2].scale(0.4)),
5359 fixed_vars[2].mul(&fixed_vars[2]).sub(&fixed_vars[0]),
5360 ];
5361 let fixed_direct = Order2::symmetric_quadratic_form(&fixed_inputs, &coefficients);
5362 let fixed_scalar = symmetric_quadratic_form_default(
5363 &fixed_inputs,
5364 &coefficients,
5365 Order2::constant,
5366 JetField::add,
5367 JetField::mul,
5368 JetField::scale,
5369 );
5370 let weights = [0.7, -1.1, 0.35];
5371 let fixed_linear_direct = Order2::linear_combination(&fixed_inputs, &weights);
5372 let fixed_linear_scalar = linear_combination_default(
5373 &fixed_inputs,
5374 &weights,
5375 Order2::constant,
5376 JetField::add,
5377 JetField::scale,
5378 );
5379 let derivative_stacks = [
5380 [0.8, -0.3, 0.7, 0.0, 0.0],
5381 [-0.2, 1.1, -0.4, 0.0, 0.0],
5382 [1.4, 0.25, 0.6, 0.0, 0.0],
5383 ];
5384 let fixed_add_direct = fixed_inputs[0].add_constant(0.65);
5385 let fixed_add_scalar = fixed_inputs[0].add(&Order2::constant(0.65));
5386 let fixed_multiply_add_direct =
5387 fixed_inputs[0].multiply_add(&fixed_inputs[1], &fixed_inputs[2]);
5388 let fixed_multiply_add_scalar = multiply_add_default(
5389 &fixed_inputs[0],
5390 &fixed_inputs[1],
5391 &fixed_inputs[2],
5392 JetField::mul,
5393 JetField::add,
5394 );
5395 let fixed_composed_direct = Order2::composed_sum(&fixed_inputs, &derivative_stacks);
5396 let fixed_composed_scalar = composed_sum_default(
5397 &fixed_inputs,
5398 &derivative_stacks,
5399 Order2::constant,
5400 JetField::add,
5401 JetField::compose_unary,
5402 );
5403
5404 let value_vars: [RuntimeValue; K] =
5405 std::array::from_fn(|axis| RuntimeValue::variable(values[axis], axis, K, &()));
5406 let value_inputs = [
5407 value_vars[0].mul(&value_vars[1]).add(&value_vars[3]),
5408 value_vars[1].exp().add(&value_vars[2].scale(0.4)),
5409 value_vars[2].mul(&value_vars[2]).sub(&value_vars[0]),
5410 ];
5411 let value_quadratic =
5412 RuntimeValue::symmetric_quadratic_form(&value_inputs, &coefficients, K, &());
5413 let value_linear = RuntimeValue::linear_combination(&value_inputs, &weights, K, &());
5414 let value_composed = RuntimeValue::composed_sum(&value_inputs, &derivative_stacks, K, &());
5415
5416 let arena = DynamicJetArena::new();
5417 let dynamic_vars: [DynamicOrder2<'_>; K] =
5418 std::array::from_fn(|axis| DynamicOrder2::variable(values[axis], axis, K, &arena));
5419 let dynamic_inputs = [
5420 dynamic_vars[0].mul(&dynamic_vars[1]).add(&dynamic_vars[3]),
5421 dynamic_vars[1].exp().add(&dynamic_vars[2].scale(0.4)),
5422 dynamic_vars[2].mul(&dynamic_vars[2]).sub(&dynamic_vars[0]),
5423 ];
5424 let dynamic_direct =
5425 DynamicOrder2::symmetric_quadratic_form(&dynamic_inputs, &coefficients, K, &arena);
5426 let dynamic_scalar = symmetric_quadratic_form_default(
5427 &dynamic_inputs,
5428 &coefficients,
5429 |value| DynamicOrder2::constant(value, K, &arena),
5430 RuntimeJetScalar::add,
5431 RuntimeJetScalar::mul,
5432 RuntimeJetScalar::scale,
5433 );
5434 let dynamic_linear_direct =
5435 DynamicOrder2::linear_combination(&dynamic_inputs, &weights, K, &arena);
5436 let dynamic_linear_scalar = linear_combination_default(
5437 &dynamic_inputs,
5438 &weights,
5439 |value| DynamicOrder2::constant(value, K, &arena),
5440 RuntimeJetScalar::add,
5441 RuntimeJetScalar::scale,
5442 );
5443 let dynamic_add_direct = dynamic_inputs[0].add_constant(0.65, &arena);
5444 let dynamic_add_scalar = dynamic_inputs[0].add(&DynamicOrder2::constant(0.65, K, &arena));
5445 let dynamic_multiply_add_direct =
5446 dynamic_inputs[0].multiply_add(&dynamic_inputs[1], &dynamic_inputs[2]);
5447 let dynamic_multiply_add_scalar = multiply_add_default(
5448 &dynamic_inputs[0],
5449 &dynamic_inputs[1],
5450 &dynamic_inputs[2],
5451 RuntimeJetScalar::mul,
5452 RuntimeJetScalar::add,
5453 );
5454 let dynamic_composed_direct =
5455 DynamicOrder2::composed_sum(&dynamic_inputs, &derivative_stacks, K, &arena);
5456 let dynamic_composed_scalar = composed_sum_default(
5457 &dynamic_inputs,
5458 &derivative_stacks,
5459 |value| DynamicOrder2::constant(value, K, &arena),
5460 RuntimeJetScalar::add,
5461 RuntimeJetScalar::compose_unary,
5462 );
5463
5464 let tolerance = 2.0e-13;
5465 for (label, actual, expected) in [
5466 ("fixed value", fixed_direct.value(), fixed_scalar.value()),
5467 (
5468 "zero-order quadratic value",
5469 value_quadratic.value(),
5470 fixed_direct.value(),
5471 ),
5472 (
5473 "zero-order linear value",
5474 value_linear.value(),
5475 fixed_linear_direct.value(),
5476 ),
5477 (
5478 "zero-order composed value",
5479 value_composed.value(),
5480 fixed_composed_direct.value(),
5481 ),
5482 (
5483 "dynamic value",
5484 dynamic_direct.value(),
5485 dynamic_scalar.value(),
5486 ),
5487 (
5488 "fixed linear value",
5489 fixed_linear_direct.value(),
5490 fixed_linear_scalar.value(),
5491 ),
5492 (
5493 "dynamic linear value",
5494 dynamic_linear_direct.value(),
5495 dynamic_linear_scalar.value(),
5496 ),
5497 (
5498 "fixed add-constant value",
5499 fixed_add_direct.value(),
5500 fixed_add_scalar.value(),
5501 ),
5502 (
5503 "dynamic add-constant value",
5504 dynamic_add_direct.value(),
5505 dynamic_add_scalar.value(),
5506 ),
5507 (
5508 "fixed multiply-add value",
5509 fixed_multiply_add_direct.value(),
5510 fixed_multiply_add_scalar.value(),
5511 ),
5512 (
5513 "dynamic multiply-add value",
5514 dynamic_multiply_add_direct.value(),
5515 dynamic_multiply_add_scalar.value(),
5516 ),
5517 (
5518 "fixed composed-sum value",
5519 fixed_composed_direct.value(),
5520 fixed_composed_scalar.value(),
5521 ),
5522 (
5523 "dynamic composed-sum value",
5524 dynamic_composed_direct.value(),
5525 dynamic_composed_scalar.value(),
5526 ),
5527 ] {
5528 assert!(
5529 (actual - expected).abs() <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5530 "{label}: direct={actual:+.16e}, scalar={expected:+.16e}"
5531 );
5532 }
5533 for primary_a in 0..K {
5534 for (label, actual, expected) in [
5535 (
5536 "fixed gradient",
5537 fixed_direct.g()[primary_a],
5538 fixed_scalar.g()[primary_a],
5539 ),
5540 (
5541 "dynamic gradient",
5542 dynamic_direct.g()[primary_a],
5543 dynamic_scalar.g()[primary_a],
5544 ),
5545 (
5546 "fixed linear gradient",
5547 fixed_linear_direct.g()[primary_a],
5548 fixed_linear_scalar.g()[primary_a],
5549 ),
5550 (
5551 "dynamic linear gradient",
5552 dynamic_linear_direct.g()[primary_a],
5553 dynamic_linear_scalar.g()[primary_a],
5554 ),
5555 (
5556 "fixed add-constant gradient",
5557 fixed_add_direct.g()[primary_a],
5558 fixed_add_scalar.g()[primary_a],
5559 ),
5560 (
5561 "dynamic add-constant gradient",
5562 dynamic_add_direct.g()[primary_a],
5563 dynamic_add_scalar.g()[primary_a],
5564 ),
5565 (
5566 "fixed multiply-add gradient",
5567 fixed_multiply_add_direct.g()[primary_a],
5568 fixed_multiply_add_scalar.g()[primary_a],
5569 ),
5570 (
5571 "dynamic multiply-add gradient",
5572 dynamic_multiply_add_direct.g()[primary_a],
5573 dynamic_multiply_add_scalar.g()[primary_a],
5574 ),
5575 (
5576 "fixed composed-sum gradient",
5577 fixed_composed_direct.g()[primary_a],
5578 fixed_composed_scalar.g()[primary_a],
5579 ),
5580 (
5581 "dynamic composed-sum gradient",
5582 dynamic_composed_direct.g()[primary_a],
5583 dynamic_composed_scalar.g()[primary_a],
5584 ),
5585 ] {
5586 assert!(
5587 (actual - expected).abs()
5588 <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5589 "{label}[{primary_a}]: direct={actual:+.16e}, scalar={expected:+.16e}"
5590 );
5591 }
5592 for primary_b in 0..K {
5593 for (label, actual, expected) in [
5594 (
5595 "fixed Hessian",
5596 fixed_direct.h()[primary_a][primary_b],
5597 fixed_scalar.h()[primary_a][primary_b],
5598 ),
5599 (
5600 "dynamic Hessian",
5601 dynamic_direct.h_at(primary_a, primary_b),
5602 dynamic_scalar.h_at(primary_a, primary_b),
5603 ),
5604 (
5605 "fixed linear Hessian",
5606 fixed_linear_direct.h()[primary_a][primary_b],
5607 fixed_linear_scalar.h()[primary_a][primary_b],
5608 ),
5609 (
5610 "dynamic linear Hessian",
5611 dynamic_linear_direct.h_at(primary_a, primary_b),
5612 dynamic_linear_scalar.h_at(primary_a, primary_b),
5613 ),
5614 (
5615 "fixed add-constant Hessian",
5616 fixed_add_direct.h()[primary_a][primary_b],
5617 fixed_add_scalar.h()[primary_a][primary_b],
5618 ),
5619 (
5620 "dynamic add-constant Hessian",
5621 dynamic_add_direct.h_at(primary_a, primary_b),
5622 dynamic_add_scalar.h_at(primary_a, primary_b),
5623 ),
5624 (
5625 "fixed multiply-add Hessian",
5626 fixed_multiply_add_direct.h()[primary_a][primary_b],
5627 fixed_multiply_add_scalar.h()[primary_a][primary_b],
5628 ),
5629 (
5630 "dynamic multiply-add Hessian",
5631 dynamic_multiply_add_direct.h_at(primary_a, primary_b),
5632 dynamic_multiply_add_scalar.h_at(primary_a, primary_b),
5633 ),
5634 (
5635 "fixed composed-sum Hessian",
5636 fixed_composed_direct.h()[primary_a][primary_b],
5637 fixed_composed_scalar.h()[primary_a][primary_b],
5638 ),
5639 (
5640 "dynamic composed-sum Hessian",
5641 dynamic_composed_direct.h_at(primary_a, primary_b),
5642 dynamic_composed_scalar.h_at(primary_a, primary_b),
5643 ),
5644 ] {
5645 assert!(
5646 (actual - expected).abs()
5647 <= tolerance * actual.abs().max(expected.abs()).max(1.0),
5648 "{label}[{primary_a},{primary_b}]: direct={actual:+.16e}, scalar={expected:+.16e}"
5649 );
5650 }
5651 }
5652 }
5653 }
5654
5655 #[test]
5656 fn compiled_product_affine_and_fused_nodes_match_scalar_program_randomized() {
5657 const K: usize = 4;
5658 const TERMS: usize = 10;
5659
5660 fn sample(state: &mut u64) -> f64 {
5661 *state ^= *state << 13;
5662 *state ^= *state >> 7;
5663 *state ^= *state << 17;
5664 let unit = (*state >> 11) as f64 * (1.0 / ((1_u64 << 53) as f64));
5665 2.0 * unit - 1.0
5666 }
5667
5668 fn arbitrary_order2<const K: usize>(state: &mut u64) -> Order2<K> {
5669 let mut tower = crate::jet_tower::Tower2::zero();
5670 tower.v = sample(state);
5671 for primary in 0..K {
5672 tower.g[primary] = sample(state);
5673 for other in primary..K {
5674 let channel = sample(state);
5675 tower.h[primary][other] = channel;
5676 tower.h[other][primary] = channel;
5677 }
5678 }
5679 Order2(tower)
5680 }
5681
5682 fn close(actual: f64, expected: f64, case: usize, label: &str) {
5683 let tolerance = 2.0e-12 * actual.abs().max(expected.abs()).max(1.0);
5684 assert!(
5685 (actual - expected).abs() <= tolerance,
5686 "case {case} {label}: direct={actual:+.16e}, scalar={expected:+.16e}, tolerance={tolerance:.3e}"
5687 );
5688 }
5689
5690 let mut state = 0x932a_ff1e_c0de_5eed_u64;
5691 for case in 0..256 {
5692 let fixed_inputs: [Order2<K>; TERMS] =
5696 std::array::from_fn(|_| arbitrary_order2(&mut state));
5697 let mut input_scales: [f64; TERMS] = std::array::from_fn(|_| sample(&mut state));
5698 input_scales[0] = -1.25;
5699 input_scales[1] = 0.0;
5700 input_scales[4] = 1.25;
5701 let addend_scales: [f64; TERMS] = std::array::from_fn(|term| match term % 4 {
5702 0 => 0.0,
5703 1 => 1.0,
5704 2 => -0.75,
5705 _ => 0.35,
5706 });
5707 let derivative_stacks: [[f64; 5]; TERMS] =
5708 std::array::from_fn(|_| std::array::from_fn(|_| sample(&mut state)));
5709 let input_shift = sample(&mut state);
5710 let mut fixed_lefts = std::array::from_fn(|term| &fixed_inputs[term]);
5711 fixed_lefts[4] = &fixed_inputs[0];
5712 fixed_lefts[5] = &fixed_inputs[0];
5713 let fixed_right = &fixed_inputs[1];
5714 let fixed_addend = &fixed_inputs[2];
5715
5716 let fixed_product_direct = fixed_inputs[0].product(&fixed_inputs[1]);
5717 let fixed_product_scalar = fixed_inputs[0].mul(&fixed_inputs[1]);
5718 let fixed_affine_direct =
5719 fixed_inputs[2].affine_compose(input_scales[2], input_shift, derivative_stacks[2]);
5720 let fixed_affine_scalar = affine_compose_default(
5721 &fixed_inputs[2],
5722 input_scales[2],
5723 input_shift,
5724 derivative_stacks[2],
5725 JetField::scale,
5726 Order2::add_constant,
5727 JetField::compose_unary,
5728 );
5729 let fixed_sum_direct =
5730 Order2::affine_composed_sum(&fixed_inputs, &input_scales, &derivative_stacks);
5731 let fixed_sum_scalar = affine_composed_sum_default(
5732 &fixed_inputs,
5733 &input_scales,
5734 &derivative_stacks,
5735 Order2::constant,
5736 JetField::add,
5737 JetField::scale,
5738 Order2::add_constant,
5739 JetField::compose_unary,
5740 );
5741 let fixed_fused_direct = Order2::shared_multiply_add_affine_composed_sum(
5742 &fixed_lefts,
5743 fixed_right,
5744 fixed_addend,
5745 &addend_scales,
5746 &input_scales,
5747 &derivative_stacks,
5748 );
5749 let fixed_fused_scalar = shared_multiply_add_affine_composed_sum_default(
5750 &fixed_lefts,
5751 fixed_right,
5752 fixed_addend,
5753 &addend_scales,
5754 &input_scales,
5755 &derivative_stacks,
5756 Order2::constant,
5757 JetField::add,
5758 JetField::mul,
5759 JetField::scale,
5760 Order2::multiply_add,
5761 Order2::affine_compose,
5762 );
5763
5764 let arena = DynamicJetArena::new();
5765 let dynamic_inputs: [DynamicOrder2<'_>; TERMS] = std::array::from_fn(|term| {
5766 DynamicOrder2::from_channel_functions(
5767 fixed_inputs[term].value(),
5768 K,
5769 &arena,
5770 |primary| fixed_inputs[term].g()[primary],
5771 |primary, other| fixed_inputs[term].h()[primary][other],
5772 )
5773 });
5774 let dynamic_product_direct = dynamic_inputs[0].product(&dynamic_inputs[1]);
5775 let dynamic_product_scalar = dynamic_inputs[0].mul(&dynamic_inputs[1]);
5776 let dynamic_affine_direct = dynamic_inputs[2].affine_compose(
5777 input_scales[2],
5778 input_shift,
5779 derivative_stacks[2],
5780 &arena,
5781 );
5782 let dynamic_affine_scalar = affine_compose_default(
5783 &dynamic_inputs[2],
5784 input_scales[2],
5785 input_shift,
5786 derivative_stacks[2],
5787 RuntimeJetScalar::scale,
5788 |input, constant| input.add_constant(constant, &arena),
5789 RuntimeJetScalar::compose_unary,
5790 );
5791 let dynamic_sum_direct = DynamicOrder2::affine_composed_sum(
5792 &dynamic_inputs,
5793 &input_scales,
5794 &derivative_stacks,
5795 K,
5796 &arena,
5797 );
5798 let dynamic_sum_scalar = affine_composed_sum_default(
5799 &dynamic_inputs,
5800 &input_scales,
5801 &derivative_stacks,
5802 |value| DynamicOrder2::constant(value, K, &arena),
5803 RuntimeJetScalar::add,
5804 RuntimeJetScalar::scale,
5805 |input, constant| input.add_constant(constant, &arena),
5806 RuntimeJetScalar::compose_unary,
5807 );
5808 let mut dynamic_lefts = std::array::from_fn(|term| &dynamic_inputs[term]);
5809 dynamic_lefts[4] = &dynamic_inputs[0];
5810 dynamic_lefts[5] = &dynamic_inputs[0];
5811 let dynamic_right = &dynamic_inputs[1];
5812 let dynamic_addend = &dynamic_inputs[2];
5813 let dynamic_fused_direct = DynamicOrder2::shared_multiply_add_affine_composed_sum(
5814 &dynamic_lefts,
5815 dynamic_right,
5816 dynamic_addend,
5817 &addend_scales,
5818 &input_scales,
5819 &derivative_stacks,
5820 K,
5821 &arena,
5822 );
5823 let dynamic_fused_scalar = shared_multiply_add_affine_composed_sum_default(
5824 &dynamic_lefts,
5825 dynamic_right,
5826 dynamic_addend,
5827 &addend_scales,
5828 &input_scales,
5829 &derivative_stacks,
5830 |value| DynamicOrder2::constant(value, K, &arena),
5831 RuntimeJetScalar::add,
5832 RuntimeJetScalar::mul,
5833 RuntimeJetScalar::scale,
5834 RuntimeJetScalar::multiply_add,
5835 |input, scale, shift, stack| input.affine_compose(scale, shift, stack, &arena),
5836 );
5837
5838 for (label, actual, expected) in [
5839 (
5840 "fixed product value",
5841 fixed_product_direct.value(),
5842 fixed_product_scalar.value(),
5843 ),
5844 (
5845 "fixed affine value",
5846 fixed_affine_direct.value(),
5847 fixed_affine_scalar.value(),
5848 ),
5849 (
5850 "fixed affine sum value",
5851 fixed_sum_direct.value(),
5852 fixed_sum_scalar.value(),
5853 ),
5854 (
5855 "fixed fused value",
5856 fixed_fused_direct.value(),
5857 fixed_fused_scalar.value(),
5858 ),
5859 (
5860 "dynamic product value",
5861 dynamic_product_direct.value(),
5862 dynamic_product_scalar.value(),
5863 ),
5864 (
5865 "dynamic affine value",
5866 dynamic_affine_direct.value(),
5867 dynamic_affine_scalar.value(),
5868 ),
5869 (
5870 "dynamic affine sum value",
5871 dynamic_sum_direct.value(),
5872 dynamic_sum_scalar.value(),
5873 ),
5874 (
5875 "dynamic fused value",
5876 dynamic_fused_direct.value(),
5877 dynamic_fused_scalar.value(),
5878 ),
5879 ] {
5880 close(actual, expected, case, label);
5881 }
5882 for primary in 0..K {
5883 for (label, actual, expected) in [
5884 (
5885 "fixed product gradient",
5886 fixed_product_direct.g()[primary],
5887 fixed_product_scalar.g()[primary],
5888 ),
5889 (
5890 "fixed affine gradient",
5891 fixed_affine_direct.g()[primary],
5892 fixed_affine_scalar.g()[primary],
5893 ),
5894 (
5895 "fixed affine sum gradient",
5896 fixed_sum_direct.g()[primary],
5897 fixed_sum_scalar.g()[primary],
5898 ),
5899 (
5900 "fixed fused gradient",
5901 fixed_fused_direct.g()[primary],
5902 fixed_fused_scalar.g()[primary],
5903 ),
5904 (
5905 "dynamic product gradient",
5906 dynamic_product_direct.g()[primary],
5907 dynamic_product_scalar.g()[primary],
5908 ),
5909 (
5910 "dynamic affine gradient",
5911 dynamic_affine_direct.g()[primary],
5912 dynamic_affine_scalar.g()[primary],
5913 ),
5914 (
5915 "dynamic affine sum gradient",
5916 dynamic_sum_direct.g()[primary],
5917 dynamic_sum_scalar.g()[primary],
5918 ),
5919 (
5920 "dynamic fused gradient",
5921 dynamic_fused_direct.g()[primary],
5922 dynamic_fused_scalar.g()[primary],
5923 ),
5924 ] {
5925 close(actual, expected, case, label);
5926 }
5927 for other in 0..K {
5928 for (label, actual, expected) in [
5929 (
5930 "fixed product Hessian",
5931 fixed_product_direct.h()[primary][other],
5932 fixed_product_scalar.h()[primary][other],
5933 ),
5934 (
5935 "fixed affine Hessian",
5936 fixed_affine_direct.h()[primary][other],
5937 fixed_affine_scalar.h()[primary][other],
5938 ),
5939 (
5940 "fixed affine sum Hessian",
5941 fixed_sum_direct.h()[primary][other],
5942 fixed_sum_scalar.h()[primary][other],
5943 ),
5944 (
5945 "fixed fused Hessian",
5946 fixed_fused_direct.h()[primary][other],
5947 fixed_fused_scalar.h()[primary][other],
5948 ),
5949 (
5950 "dynamic product Hessian",
5951 dynamic_product_direct.h_at(primary, other),
5952 dynamic_product_scalar.h_at(primary, other),
5953 ),
5954 (
5955 "dynamic affine Hessian",
5956 dynamic_affine_direct.h_at(primary, other),
5957 dynamic_affine_scalar.h_at(primary, other),
5958 ),
5959 (
5960 "dynamic affine sum Hessian",
5961 dynamic_sum_direct.h_at(primary, other),
5962 dynamic_sum_scalar.h_at(primary, other),
5963 ),
5964 (
5965 "dynamic fused Hessian",
5966 dynamic_fused_direct.h_at(primary, other),
5967 dynamic_fused_scalar.h_at(primary, other),
5968 ),
5969 ] {
5970 close(actual, expected, case, label);
5971 }
5972 }
5973 }
5974 }
5975 }
5976
5977 #[test]
5978 fn shared_product_composition_accepts_empty_expression() {
5979 const K: usize = 4;
5980 let fixed_terms: [&Order2<K>; 0] = [];
5981 let fixed_shared = Order2::constant(1.0);
5982 let scales: [f64; 0] = [];
5983 let stacks: [[f64; 5]; 0] = [];
5984 let fixed = Order2::shared_multiply_add_affine_composed_sum(
5985 &fixed_terms,
5986 &fixed_shared,
5987 &fixed_shared,
5988 &scales,
5989 &scales,
5990 &stacks,
5991 );
5992 assert_eq!(fixed.value().to_bits(), 0.0_f64.to_bits());
5993 assert!(fixed.g().iter().all(|&channel| channel == 0.0));
5994 assert!(fixed.h().iter().flatten().all(|&channel| channel == 0.0));
5995
5996 let value_terms: [&RuntimeValue; 0] = [];
5997 let value_shared = RuntimeValue::constant(1.0, K, &());
5998 let value = RuntimeValue::shared_multiply_add_affine_composed_sum(
5999 &value_terms,
6000 &value_shared,
6001 &value_shared,
6002 &scales,
6003 &scales,
6004 &stacks,
6005 K,
6006 &(),
6007 );
6008 assert_eq!(value.value().to_bits(), 0.0_f64.to_bits());
6009 assert_eq!(value.dimension(), K);
6010
6011 let arena = DynamicJetArena::new();
6012 let dynamic_terms: [&DynamicOrder2<'_>; 0] = [];
6013 let dynamic_shared = DynamicOrder2::constant(1.0, K, &arena);
6014 let dynamic = DynamicOrder2::shared_multiply_add_affine_composed_sum(
6015 &dynamic_terms,
6016 &dynamic_shared,
6017 &dynamic_shared,
6018 &scales,
6019 &scales,
6020 &stacks,
6021 K,
6022 &arena,
6023 );
6024 assert_eq!(dynamic.value().to_bits(), 0.0_f64.to_bits());
6025 assert!((0..K).all(|axis| dynamic.g()[axis] == 0.0));
6026 assert!((0..K).all(|row| (0..K).all(|column| dynamic.h_at(row, column) == 0.0)));
6027 }
6028
6029 #[test]
6030 fn runtime_fused_product_composition_preserves_tower4_channels() {
6031 const K: usize = 2;
6032 const N: usize = 9;
6033 let vars = [
6034 Tower4::<K>::variable(0.37, 0),
6035 Tower4::<K>::variable(-0.61, 1),
6036 ];
6037 let upstream = [
6038 vars[0].mul(&vars[1]).add(&vars[0].exp()),
6039 vars[1].mul(&vars[1]).add(&vars[0].scale(0.3)),
6040 vars[0].mul(&vars[0]).sub(&vars[1].scale(-0.2)),
6041 ];
6042 let mut lefts: [Tower4<K>; N] = std::array::from_fn(|term| upstream[term % upstream.len()]);
6043 lefts[4] = lefts[0];
6044 lefts[5] = lefts[0];
6045 let right = upstream[1];
6046 let addend = upstream[2];
6047 let addend_scales: [f64; N] = std::array::from_fn(|term| [-0.0, 1.0, -0.7, 0.25][term % 4]);
6048 let mut input_scales: [f64; N] =
6049 std::array::from_fn(|term| [0.0, -1.3, 0.45, 1.1][term % 4]);
6050 input_scales[0] = -1.1;
6051 input_scales[4] = 1.1;
6052 let stacks: [[f64; 5]; N] = std::array::from_fn(|term| {
6053 let t = term as f64 + 1.0;
6054 [0.17 * t, -0.11 * t, 0.07 * t, -0.03 * t, 0.013 * t]
6055 });
6056
6057 let expected = (0..N).fold(Tower4::<K>::constant(0.0), |sum, term| {
6058 let inner = if addend_scales[term] == 0.0 {
6059 lefts[term].mul(&right)
6060 } else if addend_scales[term] == 1.0 {
6061 JetScalar::multiply_add(&lefts[term], &right, &addend)
6062 } else {
6063 JetScalar::multiply_add(&lefts[term], &right, &addend.scale(addend_scales[term]))
6064 };
6065 sum.add(&JetScalar::affine_compose(
6066 &inner,
6067 input_scales[term],
6068 0.0,
6069 stacks[term],
6070 ))
6071 });
6072 let wrapped_lefts: [FixedRuntimeJet<Tower4<K>, K>; N] =
6073 std::array::from_fn(|term| FixedRuntimeJet::from_inner(lefts[term]));
6074 let wrapped_right = FixedRuntimeJet::from_inner(right);
6075 let wrapped_addend = FixedRuntimeJet::from_inner(addend);
6076 let mut wrapped_left_refs: [&FixedRuntimeJet<Tower4<K>, K>; N] =
6077 std::array::from_fn(|term| &wrapped_lefts[term]);
6078 wrapped_left_refs[4] = &wrapped_lefts[0];
6079 wrapped_left_refs[5] = &wrapped_lefts[0];
6080 let actual = FixedRuntimeJet::<Tower4<K>, K>::shared_multiply_add_affine_composed_sum(
6081 &wrapped_left_refs,
6082 &wrapped_right,
6083 &wrapped_addend,
6084 &addend_scales,
6085 &input_scales,
6086 &stacks,
6087 K,
6088 &(),
6089 )
6090 .into_inner();
6091
6092 let same = |label: &str, got: f64, want: f64| {
6093 let tolerance = 2.0e-13 * got.abs().max(want.abs()).max(1.0);
6094 assert!(
6095 (got - want).abs() <= tolerance,
6096 "{label}: got={got:+.17e}, want={want:+.17e}, tolerance={tolerance:.3e}"
6097 );
6098 };
6099 same("value", actual.v, expected.v);
6100 for a in 0..K {
6101 same("gradient", actual.g[a], expected.g[a]);
6102 for b in 0..K {
6103 same("Hessian", actual.h[a][b], expected.h[a][b]);
6104 for c in 0..K {
6105 same("third", actual.t3[a][b][c], expected.t3[a][b][c]);
6106 for d in 0..K {
6107 same("fourth", actual.t4[a][b][c][d], expected.t4[a][b][c][d]);
6108 }
6109 }
6110 }
6111 }
6112 }
6113
6114 fn row_expr<S: JetScalar<2>>(p: &[S; 2]) -> S {
6119 let g = p[0].mul(&p[1]).exp();
6120 let inner = g.add(&S::constant(2.0));
6121 let radic = p[0].mul(&p[0]).add(&S::constant(1.0)).sqrt();
6122 inner.mul(&radic).sub(&p[1].mul(&p[1]).scale(0.5))
6123 }
6124
6125 struct ExprProgram {
6127 p: [f64; 2],
6128 }
6129 impl RowProgram<2> for ExprProgram {
6130 fn n_rows(&self) -> usize {
6131 1
6132 }
6133 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
6134 if row >= self.n_rows() {
6135 return Err(format!("ExprProgram: row {row} out of range"));
6136 }
6137 Ok(self.p)
6138 }
6139 fn eval<S: JetScalar<2>>(&self, row: usize, p: &[S; 2]) -> Result<S, String> {
6140 if row >= self.n_rows() {
6141 return Err(format!("ExprProgram: row {row} out of range"));
6142 }
6143 Ok(row_expr(p))
6144 }
6145 }
6146
6147 const SEED: [f64; 2] = [0.37, -0.81];
6148 const U: [f64; 2] = [0.6, -0.2];
6149 const V: [f64; 2] = [-0.4, 1.1];
6150 const TOL: f64 = 1e-10;
6151
6152 fn close(a: f64, b: f64, label: &str) {
6153 let band = TOL + TOL * a.abs().max(b.abs());
6154 assert!(
6155 (a - b).abs() <= band,
6156 "{label}: {a:+.15e} vs {b:+.15e} (band {band:.3e})"
6157 );
6158 }
6159
6160 fn tower() -> Tower4<2> {
6161 *program_full_tower(&ExprProgram { p: SEED }, 0).expect("tower")
6162 }
6163
6164 #[test]
6166 fn order2_matches_tower_value_grad_hessian() {
6167 let t = tower();
6168 let vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6169 let s = row_expr(&vars);
6170 close(s.value(), t.v, "value");
6171 for a in 0..2 {
6172 close(s.0.g[a], t.g[a], &format!("grad[{a}]"));
6173 for b in 0..2 {
6174 close(s.h()[a][b], t.h[a][b], &format!("hess[{a}][{b}]"));
6175 }
6176 }
6177 }
6178
6179 #[test]
6180 fn mapped_order2_accumulator_matches_dense_overlapping_atoms() {
6181 const K: usize = 4;
6182 let p = [0.2_f64, 0.7, -0.4, 0.3];
6183 let dense_vars: [Order2<K>; K] =
6184 std::array::from_fn(|axis| Order2::variable(p[axis], axis));
6185 let dense_q0 = dense_vars[3].mul(&dense_vars[1]).add(&dense_vars[3].exp());
6186 let dense_q1 = dense_vars[1].mul(&dense_vars[2]).sub(&dense_vars[2]);
6187 let dense = dense_q0.ln().add(&dense_q1.exp());
6188
6189 let local_q0_vars: [Order2<2>; 2] =
6190 std::array::from_fn(|axis| Order2::variable(p[[3, 1][axis]], axis));
6191 let local_q0 = local_q0_vars[0]
6192 .mul(&local_q0_vars[1])
6193 .add(&local_q0_vars[0].exp());
6194 let local_q1_vars: [Order2<2>; 2] =
6195 std::array::from_fn(|axis| Order2::variable(p[[1, 2][axis]], axis));
6196 let local_q1 = local_q1_vars[0]
6197 .mul(&local_q1_vars[1])
6198 .sub(&local_q1_vars[1]);
6199
6200 let q0 = local_q0.value();
6201 let q1_exp = local_q1.value().exp();
6202 let mut lowered = MappedOrder2Accumulator::<K>::zero();
6203 lowered.add_composed(
6204 &local_q0,
6205 [3, 1],
6206 [q0.ln(), q0.recip(), -1.0 / (q0 * q0)],
6207 false,
6208 [false, false],
6209 [false, false, false],
6210 );
6211 lowered.add_composed(
6212 &local_q1,
6213 [1, 2],
6214 [q1_exp, q1_exp, q1_exp],
6215 true,
6216 [true, false],
6217 [true, false, false],
6218 );
6219 let (value, gradient, hessian) = lowered.into_channels();
6220
6221 close(value, dense.value(), "mapped value");
6222 for i in 0..K {
6223 close(gradient[i], dense.g()[i], &format!("mapped gradient[{i}]"));
6224 for j in 0..K {
6225 close(
6226 hessian[i][j],
6227 dense.h()[i][j],
6228 &format!("mapped Hessian[{i},{j}]"),
6229 );
6230 }
6231 }
6232 }
6233
6234 #[test]
6235 #[should_panic(expected = "mapped atom axes must be injective")]
6236 fn mapped_order2_accumulator_rejects_duplicate_axes() {
6237 let vars: [Order2<2>; 2] = std::array::from_fn(|axis| Order2::variable(0.2, axis));
6238 let atom = vars[0].add(&vars[1]);
6239 let mut lowered = MappedOrder2Accumulator::<2>::zero();
6240 lowered.add_composed(
6241 &atom,
6242 [1, 1],
6243 [0.4, 1.0, 0.0],
6244 false,
6245 [false, false],
6246 [false, false, false],
6247 );
6248 }
6249
6250 #[test]
6251 #[should_panic(expected = "mapped atom axis must be within")]
6252 fn mapped_order2_accumulator_rejects_out_of_range_axes() {
6253 let atom = Order2::<1>::variable(0.2, 0);
6254 let mut lowered = MappedOrder2Accumulator::<2>::zero();
6255 lowered.add_composed(&atom, [2], [0.2, 1.0, 0.0], false, [false], [false]);
6256 }
6257
6258 #[test]
6259 fn dynamic_order2_accumulator_matches_dense_composed_sum() {
6260 const K: usize = 4;
6261
6262 struct Term {
6263 first: f64,
6264 second: f64,
6265 gradient: [f64; K],
6266 hessian: [[f64; K]; K],
6267 }
6268
6269 impl DynamicOrder2Term for Term {
6270 fn outer_first(&self) -> f64 {
6271 self.first
6272 }
6273
6274 fn outer_second(&self) -> f64 {
6275 self.second
6276 }
6277
6278 fn inner_gradient(&self, axis: usize) -> f64 {
6279 self.gradient[axis]
6280 }
6281
6282 fn inner_hessian(&self, row: usize, column: usize) -> f64 {
6283 self.hessian[row][column]
6284 }
6285 }
6286
6287 let p = [0.7, -0.3, 0.2, 0.8];
6288 let vars: [Order2<K>; K] = std::array::from_fn(|axis| Order2::variable(p[axis], axis));
6289 let first_atom = vars[0]
6290 .mul(&vars[1])
6291 .add(&vars[2].exp())
6292 .add(&Order2::constant(1.5));
6293 let second_atom = vars[1].mul(&vars[3]).sub(&vars[0]);
6294 let first_value = first_atom.value();
6295 let second_exp = second_atom.value().exp();
6296 let first_stack = [
6297 first_value.ln(),
6298 first_value.recip(),
6299 -1.0 / (first_value * first_value),
6300 0.0,
6301 0.0,
6302 ];
6303 let second_stack = [second_exp, second_exp, second_exp, second_exp, second_exp];
6304 let dense = first_atom
6305 .compose_unary(first_stack)
6306 .add(&second_atom.compose_unary(second_stack));
6307 let terms = [
6308 Term {
6309 first: first_stack[1],
6310 second: first_stack[2],
6311 gradient: *first_atom.g(),
6312 hessian: *first_atom.h(),
6313 },
6314 Term {
6315 first: second_stack[1],
6316 second: second_stack[2],
6317 gradient: *second_atom.g(),
6318 hessian: *second_atom.h(),
6319 },
6320 ];
6321 let (value, gradient, hessian) = DynamicOrder2Accumulator::from_composed_sum(
6322 K,
6323 first_stack[0] + second_stack[0],
6324 &terms,
6325 )
6326 .into_channels();
6327
6328 close(value, dense.value(), "dynamic value");
6329 for row in 0..K {
6330 close(
6331 gradient[row],
6332 dense.g()[row],
6333 &format!("dynamic gradient[{row}]"),
6334 );
6335 for column in 0..K {
6336 close(
6337 hessian[row * K + column],
6338 dense.h()[row][column],
6339 &format!("dynamic Hessian[{row},{column}]"),
6340 );
6341 }
6342 }
6343 }
6344
6345 #[derive(Clone, Copy, Debug)]
6346 struct FullTwoPattern;
6347
6348 impl HessianPattern<2, 3> for FullTwoPattern {
6349 const PAIRS: [(usize, usize); 3] = [(0, 0), (0, 1), (1, 1)];
6350 const PAIR_BITS: [[u128; 2]; 2] = hessian_pair_bits(Self::PAIRS);
6351 }
6352
6353 #[test]
6356 fn patterned_order2_matches_dense_order2() {
6357 type Sparse = PatternedOrder2<FullTwoPattern, 2, 3>;
6358 let dense_vars: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6359 let sparse_vars: [Sparse; 2] = std::array::from_fn(|a| Sparse::variable(SEED[a], a));
6360 let dense = row_expr(&dense_vars);
6361 let sparse = row_expr(&sparse_vars);
6362 close(sparse.value(), dense.value(), "patterned value");
6363 for i in 0..2 {
6364 close(sparse.g()[i], dense.g()[i], &format!("patterned grad[{i}]"));
6365 for j in 0..2 {
6366 close(
6367 sparse.h()[i][j],
6368 dense.h()[i][j],
6369 &format!("patterned hess[{i}][{j}]"),
6370 );
6371 }
6372 }
6373 }
6374
6375 #[test]
6379 fn compose_unary_with_scalar_seam_bit_identical() {
6380 fn rand_unit(state: &mut u64) -> f64 {
6381 let mut x = *state;
6382 x ^= x << 13;
6383 x ^= x >> 7;
6384 x ^= x << 17;
6385 *state = x;
6386 2.0 * ((x >> 11) as f64 / ((1u64 << 53) as f64)) - 1.0
6387 }
6388 fn stack(u: f64) -> [f64; 5] {
6390 [
6391 u.sin(),
6392 u.cos(),
6393 (2.0 * u).sin(),
6394 (0.5 * u).cos(),
6395 u * u - 0.3,
6396 ]
6397 }
6398 fn run<const K: usize>(state: &mut u64, n: usize) -> usize {
6399 for _ in 0..n {
6400 let base = rand_unit(state);
6403 let mut s = Order2::<K>::variable(base, 0);
6404 for a in 1..K {
6405 s = crate::nested_dual::JetField::mul(
6406 &s,
6407 &Order2::<K>::variable(rand_unit(state), a),
6408 );
6409 }
6410 let with = s.compose_unary_with(stack);
6411 let explicit = s.compose_unary(stack(s.value()));
6412 assert_eq!(with.value().to_bits(), explicit.value().to_bits(), "value");
6413 for a in 0..K {
6414 assert_eq!(with.g()[a].to_bits(), explicit.g()[a].to_bits(), "g[{a}]");
6415 for b in 0..K {
6416 assert_eq!(
6417 with.h()[a][b].to_bits(),
6418 explicit.h()[a][b].to_bits(),
6419 "h[{a}][{b}]"
6420 );
6421 }
6422 }
6423 }
6424 n
6425 }
6426 let mut st = 0x9e37_79b9_7f4a_7c15u64;
6427 let total = run::<2>(&mut st, 1100)
6428 + run::<3>(&mut st, 1100)
6429 + run::<4>(&mut st, 1100)
6430 + run::<9>(&mut st, 1100);
6431 assert_eq!(total, 4400);
6432 }
6433
6434 #[test]
6437 fn one_seed_matches_tower_third_contracted() {
6438 let t = tower();
6439 let truth = t.third_contracted(&U);
6440 let vars: [OneSeed<2>; 2] =
6441 std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
6442 let s = row_expr(&vars);
6443 close(s.value(), t.v, "value");
6445 for a in 0..2 {
6446 for b in 0..2 {
6447 close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
6448 }
6449 }
6450 let third = s.contracted_third();
6451 for a in 0..2 {
6452 for b in 0..2 {
6453 close(third[a][b], truth[a][b], &format!("third[{a}][{b}]"));
6454 }
6455 }
6456 }
6457
6458 #[test]
6463 fn fused_one_seed_channels_match_unfused_definition_932() {
6464 const K: usize = 8;
6465
6466 fn random_scalar(state: &mut u64) -> f64 {
6467 *state ^= *state << 13;
6468 *state ^= *state >> 7;
6469 *state ^= *state << 17;
6470 ((*state >> 11) as f64 / ((1_u64 << 53) as f64)) * 2.0 - 1.0
6471 }
6472
6473 fn random_order2<const N: usize>(state: &mut u64) -> Order2<N> {
6474 let mut tower = crate::jet_tower::Tower2::<N>::zero();
6475 tower.v = random_scalar(state);
6476 for axis in 0..N {
6477 tower.g[axis] = random_scalar(state);
6478 }
6479 for row in 0..N {
6480 for column in row..N {
6481 let channel = random_scalar(state);
6482 tower.h[row][column] = channel;
6483 tower.h[column][row] = channel;
6484 }
6485 }
6486 Order2(tower)
6487 }
6488
6489 fn assert_channels_close<const N: usize>(
6490 label: &str,
6491 actual: &OneSeed<N>,
6492 expected: &OneSeed<N>,
6493 ) {
6494 for (part_label, actual_part, expected_part, require_exact_symmetry) in [
6495 ("base", &actual.base.0, &expected.base.0, false),
6496 ("eps", &actual.eps.0, &expected.eps.0, true),
6497 ] {
6498 let check = |channel: &str, got: f64, want: f64| {
6499 let tolerance = 2.0e-14 * got.abs().max(want.abs()).max(1.0);
6500 assert!(
6501 (got - want).abs() <= tolerance,
6502 "{label} {part_label} {channel}: got={got:+.17e} want={want:+.17e}"
6503 );
6504 };
6505 check("value", actual_part.v, expected_part.v);
6506 for row in 0..N {
6507 check(
6508 &format!("gradient[{row}]"),
6509 actual_part.g[row],
6510 expected_part.g[row],
6511 );
6512 for column in 0..N {
6513 check(
6514 &format!("hessian[{row},{column}]"),
6515 actual_part.h[row][column],
6516 expected_part.h[row][column],
6517 );
6518 if require_exact_symmetry {
6519 assert_eq!(
6520 actual_part.h[row][column].to_bits(),
6521 actual_part.h[column][row].to_bits(),
6522 "{label} {part_label} Hessian symmetry at [{row},{column}]"
6523 );
6524 }
6525 }
6526 }
6527 }
6528 }
6529
6530 let mut state = 0x9320_1eed_5eed_cafe_u64;
6531 for sample in 0..256 {
6532 let left = OneSeed {
6533 base: random_order2::<K>(&mut state),
6534 eps: random_order2::<K>(&mut state),
6535 };
6536 let right = OneSeed {
6537 base: random_order2::<K>(&mut state),
6538 eps: random_order2::<K>(&mut state),
6539 };
6540
6541 let fused_product = left.mul(&right);
6542 let unfused_product = OneSeed {
6543 base: left.base.mul(&right.base),
6544 eps: left.base.mul(&right.eps).add(&left.eps.mul(&right.base)),
6545 };
6546 assert_channels_close(
6547 &format!("sample {sample} product"),
6548 &fused_product,
6549 &unfused_product,
6550 );
6551
6552 let derivatives: [f64; 5] = std::array::from_fn(|_| random_scalar(&mut state));
6553 let fused_composition = left.compose_unary(derivatives);
6554 let unfused_composition = OneSeed {
6555 base: left.base.compose_unary(derivatives),
6556 eps: left
6557 .base
6558 .compose_unary([
6559 derivatives[1],
6560 derivatives[2],
6561 derivatives[3],
6562 derivatives[4],
6563 derivatives[4],
6564 ])
6565 .mul(&left.eps),
6566 };
6567 assert_channels_close(
6568 &format!("sample {sample} composition"),
6569 &fused_composition,
6570 &unfused_composition,
6571 );
6572 }
6573 }
6574
6575 #[test]
6579 fn two_seed_matches_tower_fourth_contracted() {
6580 let t = tower();
6581 let truth4 = t.fourth_contracted(&U, &V);
6582 let truth3_u = t.third_contracted(&U);
6583 let truth3_v = t.third_contracted(&V);
6584 let vars: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
6585 let s = row_expr(&vars);
6586 close(s.value(), t.v, "value");
6587 for a in 0..2 {
6588 close(s.base.0.g[a], t.g[a], &format!("grad[{a}]"));
6589 for b in 0..2 {
6590 close(s.base.h()[a][b], t.h[a][b], &format!("base hess[{a}][{b}]"));
6591 close(
6592 s.eps.h()[a][b],
6593 truth3_u[a][b],
6594 &format!("eps third_u[{a}][{b}]"),
6595 );
6596 close(
6597 s.del.h()[a][b],
6598 truth3_v[a][b],
6599 &format!("del third_v[{a}][{b}]"),
6600 );
6601 }
6602 }
6603 let fourth = s.contracted_fourth();
6604 for a in 0..2 {
6605 for b in 0..2 {
6606 close(fourth[a][b], truth4[a][b], &format!("fourth[{a}][{b}]"));
6607 }
6608 }
6609 }
6610
6611 #[test]
6615 fn generic_program_seam_matches_tower_for_every_channel() {
6616 let t = tower();
6617 let o2: [Order2<2>; 2] = std::array::from_fn(|a| Order2::variable(SEED[a], a));
6619 let so2 = row_expr(&o2);
6620 close(so2.value(), t.v, "seam order2 value");
6621 let os: [OneSeed<2>; 2] =
6623 std::array::from_fn(|a| OneSeed::seed_direction(SEED[a], a, U[a]));
6624 let third = row_expr(&os).contracted_third();
6625 let truth3 = t.third_contracted(&U);
6626 for a in 0..2 {
6627 for b in 0..2 {
6628 close(third[a][b], truth3[a][b], &format!("seam third[{a}][{b}]"));
6629 }
6630 }
6631 let ts: [TwoSeed<2>; 2] = std::array::from_fn(|a| TwoSeed::seed(SEED[a], a, U[a], V[a]));
6633 let fourth = row_expr(&ts).contracted_fourth();
6634 let truth4 = t.fourth_contracted(&U, &V);
6635 for a in 0..2 {
6636 for b in 0..2 {
6637 close(
6638 fourth[a][b],
6639 truth4[a][b],
6640 &format!("seam fourth[{a}][{b}]"),
6641 );
6642 }
6643 }
6644 }
6645
6646 #[test]
6653 fn tower4_as_jetscalar_matches_program_tower_all_channels() {
6654 let t = tower();
6655 let vars: [Tower4<2>; 2] = std::array::from_fn(|a| Tower4::variable(SEED[a], a));
6656 let s = row_expr(&vars);
6657 close(s.v, t.v, "tower-jetscalar value");
6658 for a in 0..2 {
6659 close(s.g[a], t.g[a], &format!("tower-jetscalar grad[{a}]"));
6660 for b in 0..2 {
6661 close(
6662 s.h[a][b],
6663 t.h[a][b],
6664 &format!("tower-jetscalar hess[{a}][{b}]"),
6665 );
6666 for c in 0..2 {
6667 close(
6668 s.t3[a][b][c],
6669 t.t3[a][b][c],
6670 &format!("tower-jetscalar t3[{a}][{b}][{c}]"),
6671 );
6672 for d in 0..2 {
6673 close(
6674 s.t4[a][b][c][d],
6675 t.t4[a][b][c][d],
6676 &format!("tower-jetscalar t4[{a}][{b}][{c}][{d}]"),
6677 );
6678 }
6679 }
6680 }
6681 }
6682 }
6683
6684 #[test]
6688 fn runtime_directional_jets_match_fixed_packed_algebra_932() {
6689 fn expression<'arena, S: RuntimeJetScalar<'arena>>(vars: &[S]) -> S {
6690 let bilinear = vars[0].mul(&vars[1]);
6691 let curved = vars[2].scale(0.7).add(&vars[3].mul(&vars[3]).scale(-0.2));
6692 bilinear
6693 .add(&curved)
6694 .exp()
6695 .mul(&vars[4].compose_unary([0.4, -0.3, 0.2, -0.1, 0.05]))
6696 }
6697
6698 const K: usize = 5;
6699 let values = [0.2, -0.7, 0.4, 1.1, -0.3];
6700 let direction_u = [0.5, -0.2, 0.7, -0.4, 0.1];
6701 let direction_v = [-0.3, 0.8, 0.2, 0.6, -0.5];
6702 let close = |actual: f64, expected: f64| {
6703 let tolerance = 1.0e-13 * (1.0 + actual.abs().max(expected.abs()));
6704 assert!((actual - expected).abs() <= tolerance);
6705 };
6706
6707 let fixed_one: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
6708 .map(|axis| FixedRuntimeJet {
6709 inner: OneSeed::seed_direction(values[axis], axis, direction_u[axis]),
6710 })
6711 .collect();
6712 let arena_one = DynamicJetArena::new();
6713 let dynamic_one: Vec<DynamicOneSeed<'_>> = (0..K)
6714 .map(|axis| {
6715 DynamicOneSeed::seed_direction(values[axis], axis, direction_u[axis], K, &arena_one)
6716 })
6717 .collect();
6718 let fixed_third = expression(&fixed_one).into_inner().contracted_third();
6719 let dynamic_third = expression(&dynamic_one);
6720 for a in 0..K {
6721 for b in 0..K {
6722 assert_eq!(
6723 dynamic_third.contracted_third()[a * K + b].to_bits(),
6724 dynamic_third.contracted_third()[b * K + a].to_bits(),
6725 "arena third Hessian must be exactly symmetric at ({a},{b})"
6726 );
6727 close(
6728 dynamic_third.contracted_third()[a * K + b],
6729 fixed_third[a][b],
6730 );
6731 }
6732 }
6733
6734 let fixed_one_v: Vec<FixedRuntimeJet<OneSeed<K>, K>> = (0..K)
6735 .map(|axis| FixedRuntimeJet {
6736 inner: OneSeed::seed_direction(values[axis], axis, direction_v[axis]),
6737 })
6738 .collect();
6739 let fixed_third_v = expression(&fixed_one_v).into_inner().contracted_third();
6740 let batch_workspace = DynamicJetBatchWorkspace::new(2);
6741 let directions = [direction_u, direction_v];
6742 let batch_vars = batch_workspace.alloc_slice_fill_with(K, |axis| {
6743 DynamicOneSeedBatch::seed_directions(values[axis], axis, K, &batch_workspace, |lane| {
6744 directions[lane][axis]
6745 })
6746 });
6747 let dynamic_batch = expression(batch_vars);
6748 assert_eq!(dynamic_batch.lanes(), 2);
6749 for lane in 0..2 {
6750 let expected = if lane == 0 {
6751 &fixed_third
6752 } else {
6753 &fixed_third_v
6754 };
6755 for a in 0..K {
6756 for b in 0..K {
6757 close(
6758 dynamic_batch.contracted_third(lane)[a * K + b],
6759 expected[a][b],
6760 );
6761 }
6762 }
6763 }
6764
6765 let fixed_two: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
6766 .map(|axis| FixedRuntimeJet {
6767 inner: TwoSeed::seed(values[axis], axis, direction_u[axis], direction_v[axis]),
6768 })
6769 .collect();
6770 let arena_two = DynamicJetArena::new();
6771 let dynamic_two: Vec<DynamicTwoSeed<'_>> = (0..K)
6772 .map(|axis| {
6773 DynamicTwoSeed::seed(
6774 values[axis],
6775 axis,
6776 direction_u[axis],
6777 direction_v[axis],
6778 K,
6779 &arena_two,
6780 )
6781 })
6782 .collect();
6783 let fixed_fourth = expression(&fixed_two).into_inner().contracted_fourth();
6784 let dynamic_fourth = expression(&dynamic_two);
6785 for a in 0..K {
6786 for b in 0..K {
6787 close(
6788 dynamic_fourth.contracted_fourth()[a * K + b],
6789 fixed_fourth[a][b],
6790 );
6791 }
6792 }
6793
6794 let fixed_two_swapped: Vec<FixedRuntimeJet<TwoSeed<K>, K>> = (0..K)
6795 .map(|axis| {
6796 FixedRuntimeJet::from_inner(TwoSeed::seed(
6797 values[axis],
6798 axis,
6799 direction_v[axis],
6800 direction_u[axis],
6801 ))
6802 })
6803 .collect();
6804 let fixed_fourth_swapped = expression(&fixed_two_swapped)
6805 .into_inner()
6806 .contracted_fourth();
6807 let pair_workspace = DynamicJetBatchWorkspace::new(2);
6808 let direction_pairs = [(direction_u, direction_v), (direction_v, direction_u)];
6809 let pair_vars = pair_workspace.alloc_slice_fill_with(K, |axis| {
6810 DynamicTwoSeedBatch::seed_direction_pairs(
6811 values[axis],
6812 axis,
6813 K,
6814 &pair_workspace,
6815 |lane| (direction_pairs[lane].0[axis], direction_pairs[lane].1[axis]),
6816 )
6817 });
6818 let dynamic_pair_batch = expression(pair_vars);
6819 assert_eq!(dynamic_pair_batch.lanes(), 2);
6820 for lane in 0..2 {
6821 let expected = if lane == 0 {
6822 &fixed_fourth
6823 } else {
6824 &fixed_fourth_swapped
6825 };
6826 for a in 0..K {
6827 for b in 0..K {
6828 close(
6829 dynamic_pair_batch.contracted_fourth(lane)[a * K + b],
6830 expected[a][b],
6831 );
6832 }
6833 }
6834 }
6835 }
6836
6837 #[test]
6838 fn dynamic_jet_arena_compacts_fragmented_high_water_932() {
6839 const WORDS_PER_ALLOCATION: usize = 1 << 17;
6840 const ALLOCATIONS: usize = 6;
6841
6842 let mut arena = DynamicJetArena::new();
6843 for lane in 0..ALLOCATIONS {
6844 let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
6845 std::hint::black_box(allocation);
6846 }
6847 let fragmented_high_water = arena.allocated_bytes();
6848
6849 arena.reset();
6850 let compact_high_water = arena.allocated_bytes();
6851 assert!(
6852 compact_high_water >= fragmented_high_water,
6853 "compacted arena must retain the complete fragmented tape"
6854 );
6855
6856 for lane in 0..ALLOCATIONS {
6857 let allocation = arena.alloc_slice_fill_with(WORDS_PER_ALLOCATION, |_| lane as u64);
6858 std::hint::black_box(allocation);
6859 }
6860 assert_eq!(
6861 arena.allocated_bytes(),
6862 compact_high_water,
6863 "equal replay must fit in the compacted chunk"
6864 );
6865
6866 arena.reset();
6867 assert_eq!(
6868 arena.allocated_bytes(),
6869 compact_high_water,
6870 "stable reset must retain the compacted chunk"
6871 );
6872 }
6873}
6874
6875#[cfg(test)]
6876mod batch_tests {
6877 use super::{
6885 JetScalar, Lane, OneSeed, OneSeedBatch, OneSeedLane, Order2, Order2Batch, Order2Lane,
6886 TwoSeed, TwoSeedBatch, TwoSeedLane,
6887 };
6888 use crate::nested_dual::JetField;
6891
6892 trait RowAlg<const K: usize>: Copy {
6896 fn constant(c: f64) -> Self;
6897 fn add(&self, o: &Self) -> Self;
6898 fn sub(&self, o: &Self) -> Self;
6899 fn mul(&self, o: &Self) -> Self;
6900 fn scale(&self, s: f64) -> Self;
6901 fn exp(&self) -> Self;
6902 fn sqrt(&self) -> Self;
6903 fn recip(&self) -> Self;
6904 }
6905
6906 impl<const K: usize> RowAlg<K> for Order2<K> {
6907 fn constant(c: f64) -> Self {
6908 <Self as JetScalar<K>>::constant(c)
6909 }
6910 fn add(&self, o: &Self) -> Self {
6911 crate::nested_dual::JetField::add(self, o)
6912 }
6913 fn sub(&self, o: &Self) -> Self {
6914 crate::nested_dual::JetField::sub(self, o)
6915 }
6916 fn mul(&self, o: &Self) -> Self {
6917 crate::nested_dual::JetField::mul(self, o)
6918 }
6919 fn scale(&self, s: f64) -> Self {
6920 crate::nested_dual::JetField::scale(self, s)
6921 }
6922 fn exp(&self) -> Self {
6923 JetScalar::exp(self)
6924 }
6925 fn sqrt(&self) -> Self {
6926 JetScalar::sqrt(self)
6927 }
6928 fn recip(&self) -> Self {
6929 JetScalar::recip(self)
6930 }
6931 }
6932
6933 impl<L: Lane, const K: usize> RowAlg<K> for Order2Lane<L, K> {
6934 fn constant(c: f64) -> Self {
6935 Order2Lane::constant(L::splat(c))
6936 }
6937 fn add(&self, o: &Self) -> Self {
6938 Order2Lane::add(self, o)
6939 }
6940 fn sub(&self, o: &Self) -> Self {
6941 Order2Lane::sub(self, o)
6942 }
6943 fn mul(&self, o: &Self) -> Self {
6944 Order2Lane::mul(self, o)
6945 }
6946 fn scale(&self, s: f64) -> Self {
6947 Order2Lane::scale(self, s)
6948 }
6949 fn exp(&self) -> Self {
6950 Order2Lane::exp(self)
6951 }
6952 fn sqrt(&self) -> Self {
6953 Order2Lane::sqrt(self)
6954 }
6955 fn recip(&self) -> Self {
6956 Order2Lane::recip(self)
6957 }
6958 }
6959
6960 fn row_expr<const K: usize, A: RowAlg<K>>(p: &[A; K]) -> A {
6965 let mut s = A::constant(0.3);
6966 for a in 0..K {
6967 let b = (a + 1) % K;
6968 s = s.add(&p[a].mul(&p[b]).scale(0.1 + 0.05 * a as f64));
6969 }
6970 let e = s.exp();
6971 let r = s.mul(&s).add(&A::constant(1.0)).sqrt();
6972 let denom = e.add(&A::constant(2.0));
6973 e.mul(&r).sub(&s.scale(0.5)).mul(&denom.recip())
6974 }
6975
6976 fn rand_unit(state: &mut u64) -> f64 {
6978 let mut x = *state;
6979 x ^= x << 13;
6980 x ^= x >> 7;
6981 x ^= x << 17;
6982 *state = x;
6983 let u = (x >> 11) as f64 / ((1u64 << 53) as f64); 2.0 * u - 1.0
6985 }
6986
6987 fn check_k<const K: usize>(state: &mut u64, batches: usize) -> usize {
6990 let mut verified_rows = 0usize;
6991 for _ in 0..batches {
6992 let rows: [[f64; K]; 4] =
6994 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
6995
6996 let prod: [Order2<K>; 4] = std::array::from_fn(|r| {
6998 let p: [Order2<K>; K] = std::array::from_fn(|a| Order2::variable(rows[r][a], a));
6999 row_expr(&p)
7000 });
7001
7002 let scal: [Order2Lane<f64, K>; 4] = std::array::from_fn(|r| {
7004 let p: [Order2Lane<f64, K>; K] =
7005 std::array::from_fn(|a| Order2Lane::variable(rows[r][a], a));
7006 row_expr(&p)
7007 });
7008
7009 let pbatch: [Order2Batch<K>; K] = std::array::from_fn(|a| {
7011 let packed = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7012 Order2Batch::variable(packed, a)
7013 });
7014 let batch = row_expr(&pbatch);
7015
7016 for r in 0..4 {
7017 let g = prod[r].0;
7018 assert_eq!(scal[r].v.to_bits(), g.v.to_bits(), "K={K} scalar v");
7020 let lr = batch.lane(r).0;
7022 assert_eq!(lr.v.to_bits(), g.v.to_bits(), "K={K} batch lane {r} v");
7023 for a in 0..K {
7024 assert_eq!(
7025 scal[r].g[a].to_bits(),
7026 g.g[a].to_bits(),
7027 "K={K} scalar g[{a}]"
7028 );
7029 assert_eq!(
7030 lr.g[a].to_bits(),
7031 g.g[a].to_bits(),
7032 "K={K} batch lane {r} g[{a}]"
7033 );
7034 for b in 0..K {
7035 assert_eq!(
7036 scal[r].h[a][b].to_bits(),
7037 g.h[a][b].to_bits(),
7038 "K={K} scalar h[{a}][{b}]"
7039 );
7040 assert_eq!(
7041 lr.h[a][b].to_bits(),
7042 g.h[a][b].to_bits(),
7043 "K={K} batch lane {r} h[{a}][{b}]"
7044 );
7045 }
7046 }
7047 verified_rows += 1;
7048 }
7049 }
7050 verified_rows
7051 }
7052
7053 #[test]
7056 fn batch_lanes_bit_identical_to_scalar_per_row() {
7057 let mut state = 0x9E37_79B9_7F4A_7C15_u64;
7058 let mut verified = 0usize;
7059 verified += check_k::<2>(&mut state, 2000);
7060 verified += check_k::<3>(&mut state, 2000);
7061 verified += check_k::<4>(&mut state, 2000);
7062 verified += check_k::<9>(&mut state, 2000);
7063 assert_eq!(verified, 4 * 2000 * 4, "every batch row must be verified");
7065 }
7066
7067 impl<const K: usize> RowAlg<K> for OneSeed<K> {
7076 fn constant(c: f64) -> Self {
7077 <Self as JetScalar<K>>::constant(c)
7078 }
7079 fn add(&self, o: &Self) -> Self {
7080 crate::nested_dual::JetField::add(self, o)
7081 }
7082 fn sub(&self, o: &Self) -> Self {
7083 crate::nested_dual::JetField::sub(self, o)
7084 }
7085 fn mul(&self, o: &Self) -> Self {
7086 crate::nested_dual::JetField::mul(self, o)
7087 }
7088 fn scale(&self, s: f64) -> Self {
7089 crate::nested_dual::JetField::scale(self, s)
7090 }
7091 fn exp(&self) -> Self {
7092 JetScalar::exp(self)
7093 }
7094 fn sqrt(&self) -> Self {
7095 JetScalar::sqrt(self)
7096 }
7097 fn recip(&self) -> Self {
7098 JetScalar::recip(self)
7099 }
7100 }
7101
7102 impl<L: Lane, const K: usize> RowAlg<K> for OneSeedLane<L, K> {
7103 fn constant(c: f64) -> Self {
7104 OneSeedLane::constant(L::splat(c))
7105 }
7106 fn add(&self, o: &Self) -> Self {
7107 OneSeedLane::add(self, o)
7108 }
7109 fn sub(&self, o: &Self) -> Self {
7110 OneSeedLane::sub(self, o)
7111 }
7112 fn mul(&self, o: &Self) -> Self {
7113 OneSeedLane::mul(self, o)
7114 }
7115 fn scale(&self, s: f64) -> Self {
7116 OneSeedLane::scale(self, s)
7117 }
7118 fn exp(&self) -> Self {
7119 OneSeedLane::exp(self)
7120 }
7121 fn sqrt(&self) -> Self {
7122 OneSeedLane::sqrt(self)
7123 }
7124 fn recip(&self) -> Self {
7125 OneSeedLane::recip(self)
7126 }
7127 }
7128
7129 impl<const K: usize> RowAlg<K> for TwoSeed<K> {
7130 fn constant(c: f64) -> Self {
7131 <Self as JetScalar<K>>::constant(c)
7132 }
7133 fn add(&self, o: &Self) -> Self {
7134 crate::nested_dual::JetField::add(self, o)
7135 }
7136 fn sub(&self, o: &Self) -> Self {
7137 crate::nested_dual::JetField::sub(self, o)
7138 }
7139 fn mul(&self, o: &Self) -> Self {
7140 crate::nested_dual::JetField::mul(self, o)
7141 }
7142 fn scale(&self, s: f64) -> Self {
7143 crate::nested_dual::JetField::scale(self, s)
7144 }
7145 fn exp(&self) -> Self {
7146 JetScalar::exp(self)
7147 }
7148 fn sqrt(&self) -> Self {
7149 JetScalar::sqrt(self)
7150 }
7151 fn recip(&self) -> Self {
7152 JetScalar::recip(self)
7153 }
7154 }
7155
7156 impl<L: Lane, const K: usize> RowAlg<K> for TwoSeedLane<L, K> {
7157 fn constant(c: f64) -> Self {
7158 TwoSeedLane::constant(L::splat(c))
7159 }
7160 fn add(&self, o: &Self) -> Self {
7161 TwoSeedLane::add(self, o)
7162 }
7163 fn sub(&self, o: &Self) -> Self {
7164 TwoSeedLane::sub(self, o)
7165 }
7166 fn mul(&self, o: &Self) -> Self {
7167 TwoSeedLane::mul(self, o)
7168 }
7169 fn scale(&self, s: f64) -> Self {
7170 TwoSeedLane::scale(self, s)
7171 }
7172 fn exp(&self) -> Self {
7173 TwoSeedLane::exp(self)
7174 }
7175 fn sqrt(&self) -> Self {
7176 TwoSeedLane::sqrt(self)
7177 }
7178 fn recip(&self) -> Self {
7179 TwoSeedLane::recip(self)
7180 }
7181 }
7182
7183 fn check_oneseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
7184 let mut rows_checked = 0;
7185 for _ in 0..batches {
7186 let rows: [[f64; K]; 4] =
7187 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7188 let u: [[f64; K]; 4] =
7190 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7191
7192 let prod: [OneSeed<K>; 4] = std::array::from_fn(|r| {
7194 let p: [OneSeed<K>; K] =
7195 std::array::from_fn(|a| OneSeed::seed_direction(rows[r][a], a, u[r][a]));
7196 row_expr(&p)
7197 });
7198
7199 let scal: [OneSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
7201 let p: [OneSeedLane<f64, K>; K] =
7202 std::array::from_fn(|a| OneSeedLane::seed_direction(rows[r][a], a, u[r][a]));
7203 row_expr(&p)
7204 });
7205
7206 let pbatch: [OneSeedBatch<K>; K] = std::array::from_fn(|a| {
7208 let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7209 let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
7210 OneSeedBatch::seed_direction(val, a, uu)
7211 });
7212 let batch = row_expr(&pbatch);
7213
7214 for r in 0..4 {
7215 let want = prod[r].contracted_third();
7216 let got_scal = scal[r].contracted_third();
7217 let got_batch = batch.lane(r).contracted_third();
7218 assert_eq!(
7220 scal[r].base.v.to_bits(),
7221 prod[r].base.value().to_bits(),
7222 "OneSeed K={K} scalar value"
7223 );
7224 assert_eq!(
7225 batch.lane(r).base.value().to_bits(),
7226 prod[r].base.value().to_bits(),
7227 "OneSeed K={K} batch lane {r} value"
7228 );
7229 for a in 0..K {
7230 for b in 0..K {
7231 assert_eq!(
7232 got_scal[a][b].to_bits(),
7233 want[a][b].to_bits(),
7234 "OneSeed K={K} scalar third[{a}][{b}]"
7235 );
7236 assert_eq!(
7237 got_batch[a][b].to_bits(),
7238 want[a][b].to_bits(),
7239 "OneSeed K={K} batch lane {r} third[{a}][{b}]"
7240 );
7241 }
7242 }
7243 rows_checked += 1;
7244 }
7245 }
7246 rows_checked
7247 }
7248
7249 fn check_twoseed<const K: usize>(state: &mut u64, batches: usize) -> usize {
7250 let mut rows_checked = 0;
7251 for _ in 0..batches {
7252 let rows: [[f64; K]; 4] =
7253 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7254 let u: [[f64; K]; 4] =
7255 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7256 let v: [[f64; K]; 4] =
7257 std::array::from_fn(|_| std::array::from_fn(|_| rand_unit(state)));
7258
7259 let prod: [TwoSeed<K>; 4] = std::array::from_fn(|r| {
7260 let p: [TwoSeed<K>; K] =
7261 std::array::from_fn(|a| TwoSeed::seed(rows[r][a], a, u[r][a], v[r][a]));
7262 row_expr(&p)
7263 });
7264
7265 let scal: [TwoSeedLane<f64, K>; 4] = std::array::from_fn(|r| {
7266 let p: [TwoSeedLane<f64, K>; K] =
7267 std::array::from_fn(|a| TwoSeedLane::seed(rows[r][a], a, u[r][a], v[r][a]));
7268 row_expr(&p)
7269 });
7270
7271 let pbatch: [TwoSeedBatch<K>; K] = std::array::from_fn(|a| {
7272 let val = wide::f64x4::new([rows[0][a], rows[1][a], rows[2][a], rows[3][a]]);
7273 let uu = wide::f64x4::new([u[0][a], u[1][a], u[2][a], u[3][a]]);
7274 let vv = wide::f64x4::new([v[0][a], v[1][a], v[2][a], v[3][a]]);
7275 TwoSeedBatch::seed(val, a, uu, vv)
7276 });
7277 let batch = row_expr(&pbatch);
7278
7279 for r in 0..4 {
7280 let want = prod[r].contracted_fourth();
7281 let got_scal = scal[r].contracted_fourth();
7282 let got_batch = batch.lane(r).contracted_fourth();
7283 assert_eq!(
7284 scal[r].base.v.to_bits(),
7285 prod[r].base.value().to_bits(),
7286 "TwoSeed K={K} scalar value"
7287 );
7288 assert_eq!(
7289 batch.lane(r).base.value().to_bits(),
7290 prod[r].base.value().to_bits(),
7291 "TwoSeed K={K} batch lane {r} value"
7292 );
7293 for a in 0..K {
7294 for b in 0..K {
7295 assert_eq!(
7296 got_scal[a][b].to_bits(),
7297 want[a][b].to_bits(),
7298 "TwoSeed K={K} scalar fourth[{a}][{b}]"
7299 );
7300 assert_eq!(
7301 got_batch[a][b].to_bits(),
7302 want[a][b].to_bits(),
7303 "TwoSeed K={K} batch lane {r} fourth[{a}][{b}]"
7304 );
7305 }
7306 }
7307 rows_checked += 1;
7308 }
7309 }
7310 rows_checked
7311 }
7312
7313 #[test]
7317 fn oneseed_lanes_contracted_third_bit_identical() {
7318 let mut state = 0x1234_5678_9ABC_DEF0_u64;
7319 let batches = 2000;
7320 let rows_checked = check_oneseed::<2>(&mut state, batches)
7321 + check_oneseed::<3>(&mut state, batches)
7322 + check_oneseed::<4>(&mut state, batches)
7323 + check_oneseed::<9>(&mut state, batches);
7324 assert_eq!(rows_checked, 4 * batches * 4);
7327 }
7328
7329 #[test]
7333 fn twoseed_lanes_contracted_fourth_bit_identical() {
7334 let mut state = 0x0FED_CBA9_8765_4321_u64;
7335 let batches = 2000;
7336 let rows_checked = check_twoseed::<2>(&mut state, batches)
7337 + check_twoseed::<3>(&mut state, batches)
7338 + check_twoseed::<4>(&mut state, batches)
7339 + check_twoseed::<9>(&mut state, batches);
7340 assert_eq!(rows_checked, 4 * batches * 4);
7343 }
7344}
7345
7346#[cfg(test)]
7347mod unit_tests {
7348 use super::{
7349 DynamicJetArena, DynamicOrder2, JetScalar, OneSeed, Order1, Order2, RuntimeJetScalar,
7350 filtered_implicit_solve_scalar,
7351 };
7352 use crate::nested_dual::{Dual2, JetField};
7353
7354 fn family_program<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7358 let xy = x.mul(y);
7359 let exponential = theta.mul(&xy).exp();
7360 let theta_squared_x_squared = theta.mul(theta).mul(&x.mul(x)).scale(0.375);
7361 let theta_y_cubed = theta.mul(&y.mul(y).mul(y)).scale(-0.2);
7362 exponential
7363 .add(&theta_squared_x_squared)
7364 .add(&theta_y_cubed)
7365 }
7366
7367 fn analytic_family_first<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7368 let xy = x.mul(y);
7369 let exponential = theta.mul(&xy).exp();
7370 xy.mul(&exponential)
7371 .add(&theta.mul(&x.mul(x)).scale(0.75))
7372 .add(&y.mul(y).mul(y).scale(-0.2))
7373 }
7374
7375 fn analytic_family_second<const K: usize, S: JetScalar<K>>(x: &S, y: &S, theta: &S) -> S {
7376 let xy = x.mul(y);
7377 let exponential = theta.mul(&xy).exp();
7378 xy.mul(&xy).mul(&exponential).add(&x.mul(x).scale(0.75))
7379 }
7380
7381 fn assert_channel_close(actual: f64, expected: f64, channel: &str) {
7382 let tolerance = 256.0 * f64::EPSILON * (1.0 + actual.abs().max(expected.abs()));
7383 assert!(
7384 (actual - expected).abs() <= tolerance,
7385 "{channel}: actual={actual:.17e}, expected={expected:.17e}, tolerance={tolerance:.3e}"
7386 );
7387 }
7388
7389 fn assert_order2_channels<const K: usize>(
7390 actual: &Order2<K>,
7391 expected: &Order2<K>,
7392 prefix: &str,
7393 ) {
7394 assert_channel_close(actual.value(), expected.value(), &format!("{prefix}.value"));
7395 for a in 0..K {
7396 assert_channel_close(actual.g()[a], expected.g()[a], &format!("{prefix}.g[{a}]"));
7397 for b in 0..K {
7398 assert_channel_close(
7399 actual.h()[a][b],
7400 expected.h()[a][b],
7401 &format!("{prefix}.h[{a}][{b}]"),
7402 );
7403 }
7404 }
7405 }
7406
7407 #[test]
7410 fn dual2_order2_extracts_exact_family_value_gradient_hessian_channels() {
7411 const K: usize = 2;
7412 let x0 = 0.7;
7413 let y0 = -0.45;
7414 let theta0 = 0.6;
7415 let x = <Dual2<Order2<K>> as JetScalar<K>>::variable(x0, 0);
7416 let y = <Dual2<Order2<K>> as JetScalar<K>>::variable(y0, 1);
7417 let theta = Dual2 {
7418 v: Order2::constant(theta0),
7419 g: Order2::constant(1.0),
7420 h: Order2::constant(0.0),
7421 };
7422
7423 let actual = family_program(&x, &y, &theta);
7424 let reference_x = Order2::variable(x0, 0);
7425 let reference_y = Order2::variable(y0, 1);
7426 let reference_theta = Order2::constant(theta0);
7427 let expected_first = analytic_family_first(&reference_x, &reference_y, &reference_theta);
7428 let expected_second = analytic_family_second(&reference_x, &reference_y, &reference_theta);
7429
7430 assert_order2_channels(&actual.g, &expected_first, "family_first");
7431 assert_order2_channels(&actual.h, &expected_second, "family_second");
7432 }
7433
7434 #[test]
7437 fn dual2_oneseed_extracts_exact_family_hessian_drift() {
7438 const K: usize = 2;
7439 let x0 = 0.7;
7440 let y0 = -0.45;
7441 let theta0 = 0.6;
7442 let direction = [0.3, -0.8];
7443
7444 let mut x = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(x0, 0);
7445 let mut y = <Dual2<OneSeed<K>> as JetScalar<K>>::variable(y0, 1);
7446 x.v.eps = Order2::constant(direction[0]);
7447 y.v.eps = Order2::constant(direction[1]);
7448 let theta = Dual2 {
7449 v: OneSeed::constant(theta0),
7450 g: OneSeed::constant(1.0),
7451 h: OneSeed::constant(0.0),
7452 };
7453
7454 let actual = family_program(&x, &y, &theta);
7455 let reference_x = OneSeed::seed_direction(x0, 0, direction[0]);
7456 let reference_y = OneSeed::seed_direction(y0, 1, direction[1]);
7457 let reference_theta = OneSeed::constant(theta0);
7458 let expected = analytic_family_first(&reference_x, &reference_y, &reference_theta);
7459
7460 assert_order2_channels(&actual.g.eps, &expected.eps, "family_first_drift");
7461 }
7462
7463 #[test]
7467 fn order2_constant_has_zero_derivatives() {
7468 let s = Order2::<3>::constant(7.5);
7469 assert_eq!(s.value(), 7.5);
7470 for a in 0..3 {
7471 assert_eq!(s.g()[a], 0.0, "grad[{a}] should be zero");
7472 for b in 0..3 {
7473 assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
7474 }
7475 }
7476 }
7477
7478 #[test]
7480 fn order2_variable_has_unit_gradient_in_seeded_slot() {
7481 let x = -2.5_f64;
7482 let s = Order2::<4>::variable(x, 2);
7483 assert_eq!(s.value(), x);
7484 for a in 0..4 {
7485 let expected_g = if a == 2 { 1.0 } else { 0.0 };
7486 assert_eq!(s.g()[a], expected_g, "grad[{a}]");
7487 for b in 0..4 {
7488 assert_eq!(s.h()[a][b], 0.0, "hess[{a}][{b}] should be zero");
7489 }
7490 }
7491 }
7492
7493 #[test]
7496 fn order2_add_sub_roundtrip() {
7497 let p = Order2::<2>::variable(3.0, 0);
7498 let q = Order2::<2>::variable(2.0, 1);
7499 let pq = crate::nested_dual::JetField::add(&p, &q);
7500 assert_eq!(pq.value(), 5.0, "add value");
7502 let back = crate::nested_dual::JetField::sub(&pq, &q);
7503 for a in 0..2 {
7505 assert_eq!(back.g()[a], p.g()[a], "grad[{a}] roundtrip");
7506 }
7507 }
7508
7509 #[test]
7512 fn order2_mul_satisfies_leibniz_rule() {
7513 let pv = 3.0_f64;
7514 let qv = -2.0_f64;
7515 let p = Order2::<2>::variable(pv, 0);
7516 let q = Order2::<2>::variable(qv, 1);
7517 let pq = crate::nested_dual::JetField::mul(&p, &q);
7518 assert_eq!(pq.value(), pv * qv, "value = p·q");
7519 assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
7520 assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
7521 assert_eq!(pq.h()[0][1], 1.0, "∂²(p·q)/∂p∂q = 1");
7522 assert_eq!(pq.h()[1][0], 1.0, "∂²(p·q)/∂q∂p = 1 (symmetric)");
7523 assert_eq!(pq.h()[0][0], 0.0, "∂²(p·q)/∂p² = 0");
7524 assert_eq!(pq.h()[1][1], 0.0, "∂²(p·q)/∂q² = 0");
7525 }
7526
7527 #[test]
7529 fn order2_scale_multiplies_all_channels() {
7530 let p = Order2::<2>::variable(4.0, 0);
7531 let s = 2.5_f64;
7532 let ps = crate::nested_dual::JetField::scale(&p, s);
7533 assert_eq!(ps.value(), 4.0 * s);
7534 assert_eq!(ps.g()[0], 1.0 * s);
7535 assert_eq!(ps.g()[1], 0.0);
7536 }
7537
7538 #[test]
7541 fn order2_exp_derivative_stack_correct() {
7542 let p0 = 1.0_f64;
7543 let p = Order2::<1>::variable(p0, 0);
7544 let ep = JetScalar::exp(&p);
7545 let e = p0.exp();
7546 assert!((ep.value() - e).abs() < 1e-15, "exp value");
7547 assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p) = exp(p)");
7548 assert!((ep.h()[0][0] - e).abs() < 1e-15, "d²/dp² exp(p) = exp(p)");
7549 }
7550
7551 #[test]
7553 fn order2_ln_derivative_stack_correct() {
7554 let p0 = 2.0_f64;
7555 let p = Order2::<1>::variable(p0, 0);
7556 let lnp = JetScalar::ln(&p);
7557 assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
7558 assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
7559 assert!(
7560 (lnp.h()[0][0] - (-1.0 / (p0 * p0))).abs() < 1e-15,
7561 "d²/dp² ln(p) = -1/p²"
7562 );
7563 }
7564
7565 #[test]
7566 fn dynamic_order2_ln_uses_runtime_scalar_derivative_stack() {
7567 let p0 = 2.0_f64;
7568 let arena = DynamicJetArena::new();
7569 let p = DynamicOrder2::variable(p0, 0, 1, &arena);
7570 let lnp = RuntimeJetScalar::ln(&p);
7571 assert!((lnp.value() - p0.ln()).abs() < 1e-15, "ln value");
7572 assert!((lnp.g()[0] - 1.0 / p0).abs() < 1e-15, "d/dp ln(p) = 1/p");
7573 assert!(
7574 (lnp.h_at(0, 0) - (-1.0 / (p0 * p0))).abs() < 1e-15,
7575 "d²/dp² ln(p) = -1/p²"
7576 );
7577 }
7578
7579 #[test]
7581 fn order2_exp_ln_roundtrip_at_value() {
7582 let p0 = 0.8_f64;
7583 let p = Order2::<1>::variable(p0, 0);
7584 let roundtrip = JetScalar::ln(&JetScalar::exp(&p));
7585 assert!((roundtrip.value() - p0).abs() < 1e-14, "ln(exp(p)) ≈ p");
7586 }
7587
7588 #[test]
7592 fn order1_constant_has_zero_gradient() {
7593 let s = Order1::<3>::constant(-5.0);
7594 assert_eq!(s.value(), -5.0);
7595 for a in 0..3 {
7596 assert_eq!(s.g()[a], 0.0, "g[{a}] should be zero");
7597 }
7598 }
7599
7600 #[test]
7602 fn order1_variable_has_unit_gradient_in_seeded_slot() {
7603 let s = Order1::<3>::variable(2.0, 1);
7604 assert_eq!(s.value(), 2.0);
7605 assert_eq!(s.g()[0], 0.0);
7606 assert_eq!(s.g()[1], 1.0);
7607 assert_eq!(s.g()[2], 0.0);
7608 }
7609
7610 #[test]
7612 fn order1_mul_satisfies_product_rule() {
7613 let pv = 3.0_f64;
7614 let qv = -2.0_f64;
7615 let p = Order1::<2>::variable(pv, 0);
7616 let q = Order1::<2>::variable(qv, 1);
7617 let pq = crate::nested_dual::JetField::mul(&p, &q);
7618 assert_eq!(pq.value(), pv * qv);
7619 assert_eq!(pq.g()[0], qv, "∂(p·q)/∂p = q");
7620 assert_eq!(pq.g()[1], pv, "∂(p·q)/∂q = p");
7621 }
7622
7623 #[test]
7625 fn order1_exp_has_correct_value_and_gradient() {
7626 let p0 = 0.5_f64;
7627 let p = Order1::<2>::variable(p0, 0);
7628 let ep = JetScalar::exp(&p);
7629 let e = p0.exp();
7630 assert!((ep.value() - e).abs() < 1e-15, "exp value");
7631 assert!((ep.g()[0] - e).abs() < 1e-15, "d/dp exp(p)");
7632 assert_eq!(ep.g()[1], 0.0, "irrelevant gradient slot is zero");
7633 }
7634
7635 #[test]
7637 fn order1_and_order2_agree_on_value_and_gradient() {
7638 let p0 = 1.3_f64;
7639 let q0 = -0.7_f64;
7640 let p1 = Order1::<2>::variable(p0, 0);
7642 let q1 = Order1::<2>::variable(q0, 1);
7643 let expr1 = JetScalar::exp(&crate::nested_dual::JetField::add(
7644 &crate::nested_dual::JetField::mul(&p1, &q1),
7645 &p1,
7646 ));
7647
7648 let p2 = Order2::<2>::variable(p0, 0);
7649 let q2 = Order2::<2>::variable(q0, 1);
7650 let expr2 = JetScalar::exp(&crate::nested_dual::JetField::add(
7651 &crate::nested_dual::JetField::mul(&p2, &q2),
7652 &p2,
7653 ));
7654
7655 assert!(
7656 (expr1.value() - expr2.value()).abs() < 1e-14,
7657 "value mismatch"
7658 );
7659 for a in 0..2 {
7660 assert!(
7661 (expr1.g()[a] - expr2.g()[a]).abs() < 1e-14,
7662 "gradient[{a}] mismatch"
7663 );
7664 }
7665 }
7666
7667 #[test]
7672 fn filtered_implicit_solve_linear_constraint_gives_exact_jet() {
7673 let theta0 = 3.0_f64;
7674 let theta = Order2::<1>::variable(theta0, 0);
7675 let a = filtered_implicit_solve_scalar::<1, Order2<1>>(theta0, 1.0, 2, |a_jet| {
7677 crate::nested_dual::JetField::sub(a_jet, &theta)
7678 });
7679 assert!((a.value() - theta0).abs() < 1e-14, "value = theta0");
7680 assert!((a.g()[0] - 1.0).abs() < 1e-14, "gradient = 1");
7682 assert!(a.h()[0][0].abs() < 1e-14, "hessian = 0");
7684 }
7685
7686 #[test]
7689 fn filtered_implicit_solve_quadratic_constraint_matches_analytic_derivatives() {
7690 let theta0 = 4.0_f64;
7691 let a0 = theta0.sqrt();
7692 let inv_fa = 1.0 / (2.0 * a0);
7693 let theta = Order2::<1>::variable(theta0, 0);
7694 let a = filtered_implicit_solve_scalar::<1, Order2<1>>(a0, inv_fa, 2, |a_jet| {
7696 let aa = crate::nested_dual::JetField::mul(a_jet, a_jet);
7697 crate::nested_dual::JetField::sub(&aa, &theta)
7698 });
7699 let tol = 1e-12;
7700 assert!((a.value() - a0).abs() < tol, "value = sqrt(theta0)");
7701 let expected_g = 0.5 / a0;
7702 assert!(
7703 (a.g()[0] - expected_g).abs() < tol,
7704 "da/dtheta = 1/(2*sqrt)"
7705 );
7706 let expected_h = -0.25 / (theta0 * a0);
7707 assert!(
7708 (a.h()[0][0] - expected_h).abs() < tol,
7709 "d2a/dtheta2 = -1/(4*theta^1.5)"
7710 );
7711 }
7712}