1use std::collections::{BTreeSet, HashMap, HashSet};
26use std::sync::Arc;
27
28use super::nl_external::{EvalResult, ExternalArg, ExternalLibrary, ExternalResolver};
29use super::nl_reader::{BinOp, CmpOp, Expr, FuncallArg, UnaryOp};
30
31#[derive(Debug, Clone)]
35pub enum TapeOp {
36 Const(f64),
37 Var(usize),
38 Add(usize, usize),
39 Sub(usize, usize),
40 Mul(usize, usize),
41 Div(usize, usize),
42 Pow(usize, usize),
43 Neg(usize),
44 Abs(usize),
45 Sqrt(usize),
46 Exp(usize),
47 Log(usize),
48 Log10(usize),
49 Sin(usize),
50 Cos(usize),
51 Tan(usize),
52 Atan(usize),
53 Acos(usize),
54 Sinh(usize),
55 Cosh(usize),
56 Tanh(usize),
57 Asin(usize),
58 Acosh(usize),
59 Asinh(usize),
60 Atanh(usize),
61 Erf(usize),
69 XLogX(usize),
81 CEntropy(usize, usize),
90 Atan2(usize, usize),
93 Min(usize, usize),
98 Max(usize, usize),
101 Cmp(CmpOp, usize, usize),
105 And(usize, usize),
107 Or(usize, usize),
109 Not(usize),
111 Select(usize, usize, usize),
117 Funcall(Box<FuncallData>),
125}
126
127#[derive(Debug, Clone)]
132pub struct FuncallData {
133 pub lib: Arc<ExternalLibrary>,
134 pub name: String,
135 pub args: Vec<TapeFuncallArg>,
136}
137
138#[derive(Debug, Clone)]
142pub enum TapeFuncallArg {
143 Tape(usize),
144 Str(String),
145}
146
147#[inline]
163pub(crate) fn ln_ratio(a: f64, b: f64) -> f64 {
164 let q = a / b;
165 if q.is_finite() && q > 0.0 {
166 let t = (a - b) / b;
167 if t.abs() < 0.5 { t.ln_1p() } else { q.ln() }
168 } else {
169 a.ln() - b.ln()
170 }
171}
172
173#[inline]
185pub(crate) fn xlogx(a: f64) -> f64 {
186 if a == 0.0 { 0.0 } else { a * a.ln() }
187}
188
189#[inline]
197pub(crate) fn centropy(a: f64, b: f64) -> f64 {
198 if a == 0.0 { 0.0 } else { a * ln_ratio(a, b) }
199}
200
201#[inline]
203pub(crate) fn xlogx_d1(a: f64) -> f64 {
204 a.ln() + 1.0
205}
206
207#[inline]
212pub(crate) fn xlogx_d2(a: f64) -> f64 {
213 1.0 / a
214}
215
216#[inline]
218pub(crate) fn centropy_da(a: f64, b: f64) -> f64 {
219 ln_ratio(a, b) + 1.0
220}
221
222#[inline]
224pub(crate) fn centropy_db(a: f64, b: f64) -> f64 {
225 -(a / b)
226}
227
228#[inline]
230pub(crate) fn centropy_daa(a: f64) -> f64 {
231 1.0 / a
232}
233
234#[inline]
236pub(crate) fn centropy_dab(b: f64) -> f64 {
237 -1.0 / b
238}
239
240#[inline]
245pub(crate) fn centropy_dbb(a: f64, b: f64) -> f64 {
246 (a / b) / b
247}
248
249#[inline]
257pub(crate) fn atanh(u: f64) -> f64 {
258 libm::atanh(u)
259}
260
261#[inline]
296pub(crate) fn atan_d1(u: f64) -> f64 {
297 if u.abs() > 1e150 {
298 let r = 1.0 / u;
299 r * r
300 } else {
301 1.0 / (1.0 + u * u)
302 }
303}
304
305#[inline]
307pub(crate) fn atan_d2(u: f64) -> f64 {
308 let g = atan_d1(u);
309 -2.0 * (u * g) * g
310}
311
312#[inline]
314pub(crate) fn asin_d1(u: f64) -> f64 {
315 1.0 / ((1.0 - u) * (1.0 + u)).sqrt()
316}
317
318#[inline]
320pub(crate) fn asin_d2(u: f64) -> f64 {
321 let g = asin_d1(u);
322 (u * g) * g * g
323}
324
325#[inline]
327pub(crate) fn acos_d1(u: f64) -> f64 {
328 -asin_d1(u)
329}
330
331#[inline]
333pub(crate) fn acos_d2(u: f64) -> f64 {
334 -asin_d2(u)
335}
336
337#[inline]
340pub(crate) fn tanh_d1(u: f64) -> f64 {
341 let s = 1.0 / u.cosh();
342 s * s
343}
344
345#[inline]
348pub(crate) fn acosh_d1(u: f64) -> f64 {
349 1.0 / ((u - 1.0).sqrt() * (u + 1.0).sqrt())
350}
351
352#[inline]
354pub(crate) fn acosh_d2(u: f64) -> f64 {
355 let g = acosh_d1(u);
356 -(u * g) * g * g
357}
358
359#[inline]
361pub(crate) fn asinh_d1(u: f64) -> f64 {
362 1.0 / u.hypot(1.0)
363}
364
365#[inline]
367pub(crate) fn asinh_d2(u: f64) -> f64 {
368 let g = asinh_d1(u);
369 -(u * g) * g * g
370}
371
372#[inline]
374pub(crate) fn atanh_d1(u: f64) -> f64 {
375 1.0 / ((1.0 - u) * (1.0 + u))
376}
377
378#[inline]
380pub(crate) fn atanh_d2(u: f64) -> f64 {
381 let g = atanh_d1(u);
382 2.0 * (u * g) * g
383}
384
385#[doc(hidden)]
395#[inline]
396pub fn log_d2(u: f64) -> f64 {
397 if u.abs() > 1e150 {
398 -(1.0 / u) / u
399 } else {
400 -1.0 / (u * u)
401 }
402}
403
404#[doc(hidden)]
409#[inline]
410pub fn div_u_ln10(x: f64, u: f64) -> f64 {
411 let c = std::f64::consts::LN_10;
412 if u.abs() > 1e300 {
413 (x / u) / c
414 } else {
415 x / (u * c)
416 }
417}
418
419#[doc(hidden)]
421#[inline]
422pub fn log10_d2(u: f64) -> f64 {
423 let c = std::f64::consts::LN_10;
424 if u.abs() > 1e150 {
425 -((1.0 / u) / u) / c
426 } else {
427 -1.0 / (u * u * c)
428 }
429}
430
431#[doc(hidden)]
435#[inline]
436pub fn sqrt_d2(u: f64, sv: f64) -> f64 {
437 if u > 1e200 {
438 (-0.25 / u) / sv
439 } else {
440 -0.25 / (u * sv)
441 }
442}
443
444#[inline]
454fn atan2_frame(y: f64, x: f64) -> (f64, f64, f64, f64) {
455 let k = if y.abs().max(x.abs()) > 1e300 {
456 2.0
457 } else {
458 1.0
459 };
460 let (ys, xs) = (y / k, x / k);
461 (ys, xs, ys.hypot(xs), k)
462}
463
464#[inline]
471pub(crate) fn atan2_d1(y: f64, x: f64) -> (f64, f64) {
472 let (ys, xs, hs, k) = atan2_frame(y, x);
473 let (cy, cx) = (ys / hs, xs / hs);
474 ((cx / k) / hs, (-cy / k) / hs)
477}
478
479#[inline]
483pub(crate) fn atan2_d2(y: f64, x: f64) -> (f64, f64, f64) {
484 let (ys, xs, hs, k) = atan2_frame(y, x);
485 let (cy, cx) = (ys / hs, xs / hs);
486 let kk = k * k;
487 let q = 2.0 * cy * cx / kk / hs / hs;
488 let r = ((ys - xs) / hs) * ((ys + xs) / hs);
491 (-q, r / kk / hs / hs, q)
492}
493
494#[inline]
501pub(crate) fn erf(u: f64) -> f64 {
502 libm::erf(u)
503}
504
505#[inline]
510pub(crate) fn erf_d1(u: f64) -> f64 {
511 std::f64::consts::FRAC_2_SQRT_PI * (-u * u).exp()
512}
513
514#[inline]
525pub(crate) fn erf_d2(u: f64) -> f64 {
526 -2.0 * (u * erf_d1(u))
527}
528
529#[inline]
532fn cmp_holds(op: CmpOp, a: f64, b: f64) -> bool {
533 match op {
534 CmpOp::Lt => a < b,
535 CmpOp::Le => a <= b,
536 CmpOp::Eq => a == b,
537 CmpOp::Ge => a >= b,
538 CmpOp::Gt => a > b,
539 CmpOp::Ne => a != b,
540 }
541}
542
543fn funcall_to_ext_args<'a>(args: &'a [TapeFuncallArg], vals: &[f64]) -> Vec<ExternalArg<'a>> {
544 args.iter()
545 .map(|a| match a {
546 TapeFuncallArg::Tape(idx) => ExternalArg::Real(vals[*idx]),
547 TapeFuncallArg::Str(s) => ExternalArg::Str(s.as_str()),
548 })
549 .collect()
550}
551
552fn ext_eval_or_nan(
565 lib: &ExternalLibrary,
566 name: &str,
567 call_args: &[ExternalArg<'_>],
568 n_args: usize,
569 want_derivs: bool,
570 want_hes: bool,
571) -> EvalResult {
572 lib.eval(name, call_args, want_derivs, want_hes)
573 .unwrap_or_else(|_| EvalResult {
574 value: f64::NAN,
575 derivs: want_derivs.then(|| vec![f64::NAN; n_args]),
576 hessian: want_hes.then(|| vec![f64::NAN; n_args * (n_args + 1) / 2]),
577 })
578}
579
580#[derive(Debug, Clone)]
583pub struct Tape {
584 pub ops: Vec<TapeOp>,
585}
586
587impl Tape {
588 pub fn build(expr: &Expr) -> Self {
592 Self::build_with_externals(expr, &ExternalResolver::default())
593 }
594
595 pub fn build_with_externals(expr: &Expr, resolver: &ExternalResolver) -> Self {
600 let mut ops = Vec::new();
601 let mut cache: HashMap<*const Expr, usize> = HashMap::new();
602 build_recursive(expr, &mut ops, &mut cache, resolver);
603 Tape { ops }
604 }
605
606 pub fn forward(&self, x: &[f64]) -> Vec<f64> {
609 let mut vals: Vec<f64> = Vec::with_capacity(self.ops.len());
610 for op in &self.ops {
611 let v = match op {
612 TapeOp::Const(c) => *c,
613 TapeOp::Var(i) => x[*i],
614 TapeOp::Add(a, b) => vals[*a] + vals[*b],
615 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
616 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
617 TapeOp::Div(a, b) => vals[*a] / vals[*b],
618 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
619 TapeOp::Neg(a) => -vals[*a],
620 TapeOp::Abs(a) => vals[*a].abs(),
621 TapeOp::Sqrt(a) => vals[*a].sqrt(),
622 TapeOp::Exp(a) => vals[*a].exp(),
623 TapeOp::Log(a) => vals[*a].ln(),
624 TapeOp::Log10(a) => vals[*a].log10(),
625 TapeOp::Sin(a) => vals[*a].sin(),
626 TapeOp::Cos(a) => vals[*a].cos(),
627 TapeOp::Tan(a) => vals[*a].tan(),
628 TapeOp::Atan(a) => vals[*a].atan(),
629 TapeOp::Acos(a) => vals[*a].acos(),
630 TapeOp::Sinh(a) => vals[*a].sinh(),
631 TapeOp::Cosh(a) => vals[*a].cosh(),
632 TapeOp::Tanh(a) => vals[*a].tanh(),
633 TapeOp::Asin(a) => vals[*a].asin(),
634 TapeOp::Acosh(a) => vals[*a].acosh(),
635 TapeOp::Asinh(a) => vals[*a].asinh(),
636 TapeOp::Atanh(a) => atanh(vals[*a]),
637 TapeOp::Erf(a) => erf(vals[*a]),
638 TapeOp::XLogX(a) => xlogx(vals[*a]),
639 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
640 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
641 TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
642 TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
643 TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
644 TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
645 TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
646 TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
647 TapeOp::Select(c, t, e) => {
648 if vals[*c] != 0.0 {
649 vals[*t]
650 } else {
651 vals[*e]
652 }
653 }
654 TapeOp::Funcall(fc) => {
655 let FuncallData { lib, name, args } = fc.as_ref();
656 let call_args = funcall_to_ext_args(args, &vals);
657 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
658 res.value
659 }
660 };
661 vals.push(v);
662 }
663 vals
664 }
665
666 pub fn eval(&self, x: &[f64]) -> f64 {
667 let vals = self.forward(x);
668 *vals.last().unwrap_or(&0.0)
669 }
670
671 pub fn gradient_seed(&self, x: &[f64], seed: f64, grad: &mut [f64]) {
676 if seed == 0.0 || self.ops.is_empty() {
677 return;
678 }
679 let vals = self.forward(x);
680 self.reverse(&vals, seed, grad);
681 }
682
683 pub fn gradient_seed_into(
695 &self,
696 x: &[f64],
697 seed: f64,
698 grad: &mut [f64],
699 vals: &mut [f64],
700 adj: &mut [f64],
701 ) {
702 if seed == 0.0 || self.ops.is_empty() {
703 return;
704 }
705 debug_assert!(vals.len() >= self.ops.len());
706 self.forward_into(x, vals);
707 self.reverse_into(vals, seed, grad, adj);
708 }
709
710 fn reverse(&self, vals: &[f64], seed: f64, grad: &mut [f64]) {
711 let n = self.ops.len();
712 let mut adj = vec![0.0f64; n];
713 self.reverse_into(vals, seed, grad, &mut adj);
714 }
715
716 fn reverse_into(&self, vals: &[f64], seed: f64, grad: &mut [f64], adj: &mut [f64]) {
721 let n = self.ops.len();
722 debug_assert!(adj.len() >= n);
723 adj[..n].fill(0.0);
724 adj[n - 1] = seed;
725
726 for i in (0..n).rev() {
727 let a = adj[i];
728 if a == 0.0 {
729 continue;
730 }
731 match &self.ops[i] {
732 TapeOp::Const(_) => {}
733 TapeOp::Var(j) => {
734 grad[*j] += a;
735 }
736 TapeOp::Add(l, r) => {
737 adj[*l] += a;
738 adj[*r] += a;
739 }
740 TapeOp::Sub(l, r) => {
741 adj[*l] += a;
742 adj[*r] -= a;
743 }
744 TapeOp::Mul(l, r) => {
745 adj[*l] += a * vals[*r];
746 adj[*r] += a * vals[*l];
747 }
748 TapeOp::Div(l, r) => {
749 let rv = vals[*r];
753 adj[*l] += a / rv;
754 adj[*r] -= a * vals[i] / rv;
755 }
756 TapeOp::Pow(l, r) => {
757 let lv = vals[*l];
758 let rv = vals[*r];
759 if rv != 0.0 {
760 adj[*l] += a * rv * lv.powf(rv - 1.0);
761 }
762 if lv > 0.0 {
763 adj[*r] += a * vals[i] * lv.ln();
764 }
765 }
766 TapeOp::Neg(j) => {
767 adj[*j] -= a;
768 }
769 TapeOp::Abs(j) => {
770 if vals[*j] >= 0.0 {
771 adj[*j] += a;
772 } else {
773 adj[*j] -= a;
774 }
775 }
776 TapeOp::Sqrt(j) => {
777 let sv = vals[i];
778 if sv > 0.0 {
779 adj[*j] += a * 0.5 / sv;
780 }
781 }
782 TapeOp::Exp(j) => {
783 adj[*j] += a * vals[i];
784 }
785 TapeOp::Log(j) => {
786 adj[*j] += a / vals[*j];
787 }
788 TapeOp::Log10(j) => {
789 adj[*j] += div_u_ln10(a, vals[*j]);
790 }
791 TapeOp::Sin(j) => {
792 adj[*j] += a * vals[*j].cos();
793 }
794 TapeOp::Cos(j) => {
795 adj[*j] -= a * vals[*j].sin();
796 }
797 TapeOp::Tan(j) => {
798 let t = vals[i];
799 adj[*j] += a * (1.0 + t * t);
800 }
801 TapeOp::Atan(j) => {
802 adj[*j] += a * atan_d1(vals[*j]);
803 }
804 TapeOp::Acos(j) => {
805 adj[*j] += a * acos_d1(vals[*j]);
806 }
807 TapeOp::Sinh(j) => {
808 adj[*j] += a * vals[*j].cosh();
809 }
810 TapeOp::Cosh(j) => {
811 adj[*j] += a * vals[*j].sinh();
812 }
813 TapeOp::Tanh(j) => {
814 adj[*j] += a * tanh_d1(vals[*j]);
815 }
816 TapeOp::Asin(j) => {
817 adj[*j] += a * asin_d1(vals[*j]);
818 }
819 TapeOp::Acosh(j) => {
820 adj[*j] += a * acosh_d1(vals[*j]);
821 }
822 TapeOp::Asinh(j) => {
823 adj[*j] += a * asinh_d1(vals[*j]);
824 }
825 TapeOp::Atanh(j) => {
826 adj[*j] += a * atanh_d1(vals[*j]);
827 }
828 TapeOp::Erf(j) => {
829 adj[*j] += a * erf_d1(vals[*j]);
830 }
831 TapeOp::XLogX(j) => {
832 adj[*j] += a * xlogx_d1(vals[*j]);
833 }
834 TapeOp::CEntropy(l, r) => {
835 adj[*l] += a * centropy_da(vals[*l], vals[*r]);
836 adj[*r] += a * centropy_db(vals[*l], vals[*r]);
837 }
838 TapeOp::Atan2(l, r) => {
839 let y = vals[*l];
840 let x = vals[*r];
841 let (fa, fb) = atan2_d1(y, x);
842 adj[*l] += a * fa;
843 adj[*r] += a * fb;
844 }
845 TapeOp::Min(l, r) => {
849 if vals[*l] <= vals[*r] {
850 adj[*l] += a;
851 } else {
852 adj[*r] += a;
853 }
854 }
855 TapeOp::Max(l, r) => {
856 if vals[*l] >= vals[*r] {
857 adj[*l] += a;
858 } else {
859 adj[*r] += a;
860 }
861 }
862 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
865 TapeOp::Select(c, t, e) => {
868 if vals[*c] != 0.0 {
869 adj[*t] += a;
870 } else {
871 adj[*e] += a;
872 }
873 }
874 TapeOp::Funcall(fc) => {
875 let FuncallData { lib, name, args } = fc.as_ref();
876 let call_args = funcall_to_ext_args(args, vals);
877 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
878 let derivs = res.derivs.expect("want_derivs=true returns derivs");
879 let mut k = 0usize;
880 for arg in args {
881 if let TapeFuncallArg::Tape(idx) = arg {
882 adj[*idx] += a * derivs[k];
883 k += 1;
884 }
885 }
886 }
887 }
888 }
889 }
890
891 pub fn variables(&self) -> Vec<usize> {
893 let mut s: BTreeSet<usize> = BTreeSet::new();
894 for op in &self.ops {
895 if let TapeOp::Var(j) = op {
896 s.insert(*j);
897 }
898 }
899 s.into_iter().collect()
900 }
901
902 fn forward_tangent(&self, vals: &[f64], seed_var: usize, dot: &mut [f64]) {
907 let n = self.ops.len();
908 debug_assert_eq!(dot.len(), n);
909 for i in 0..n {
910 dot[i] = match &self.ops[i] {
911 TapeOp::Const(_) => 0.0,
912 TapeOp::Var(k) => {
913 if *k == seed_var {
914 1.0
915 } else {
916 0.0
917 }
918 }
919 TapeOp::Add(a, b) => dot[*a] + dot[*b],
920 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
921 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
922 TapeOp::Div(a, b) => {
923 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
931 }
932 TapeOp::Pow(a, b) => {
933 let u = vals[*a];
934 let r = vals[*b];
935 let du = dot[*a];
936 let dr = dot[*b];
937 let mut result = 0.0;
938 if r != 0.0 {
943 result += r * u.powf(r - 1.0) * du;
944 }
945 if u > 0.0 {
946 result += vals[i] * u.ln() * dr;
947 }
948 result
949 }
950 TapeOp::Neg(a) => -dot[*a],
951 TapeOp::Abs(a) => {
952 if vals[*a] >= 0.0 {
953 dot[*a]
954 } else {
955 -dot[*a]
956 }
957 }
958 TapeOp::Sqrt(a) => {
959 let sv = vals[i];
960 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
961 }
962 TapeOp::Exp(a) => dot[*a] * vals[i],
963 TapeOp::Log(a) => dot[*a] / vals[*a],
964 TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
965 TapeOp::Sin(a) => dot[*a] * vals[*a].cos(),
966 TapeOp::Cos(a) => -dot[*a] * vals[*a].sin(),
967 TapeOp::Tan(a) => {
968 let t = vals[i];
969 dot[*a] * (1.0 + t * t)
970 }
971 TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
972 TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
973 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
974 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
975 TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
976 TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
977 TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
978 TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
979 TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
980 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
981 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
982 TapeOp::CEntropy(a, b) => {
983 centropy_da(vals[*a], vals[*b]) * dot[*a]
984 + centropy_db(vals[*a], vals[*b]) * dot[*b]
985 }
986 TapeOp::Atan2(a, b) => {
987 let y = vals[*a];
988 let x = vals[*b];
989 let (fa, fb) = atan2_d1(y, x);
990 fa * dot[*a] + fb * dot[*b]
991 }
992 TapeOp::Min(a, b) => {
994 if vals[*a] <= vals[*b] {
995 dot[*a]
996 } else {
997 dot[*b]
998 }
999 }
1000 TapeOp::Max(a, b) => {
1001 if vals[*a] >= vals[*b] {
1002 dot[*a]
1003 } else {
1004 dot[*b]
1005 }
1006 }
1007 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
1008 TapeOp::Select(c, t, e) => {
1009 if vals[*c] != 0.0 {
1010 dot[*t]
1011 } else {
1012 dot[*e]
1013 }
1014 }
1015 TapeOp::Funcall(fc) => {
1016 let FuncallData { lib, name, args } = fc.as_ref();
1017 let call_args = funcall_to_ext_args(args, vals);
1018 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
1019 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1020 let mut acc = 0.0;
1021 let mut k = 0usize;
1022 for arg in args {
1023 if let TapeFuncallArg::Tape(idx) = arg {
1024 acc += derivs[k] * dot[*idx];
1025 k += 1;
1026 }
1027 }
1028 acc
1029 }
1030 };
1031 }
1032 }
1033
1034 pub fn forward_into(&self, x: &[f64], vals: &mut [f64]) {
1038 let n = self.ops.len();
1039 debug_assert!(vals.len() >= n);
1040 for i in 0..n {
1041 vals[i] = match &self.ops[i] {
1042 TapeOp::Const(c) => *c,
1043 TapeOp::Var(j) => x[*j],
1044 TapeOp::Add(a, b) => vals[*a] + vals[*b],
1045 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
1046 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
1047 TapeOp::Div(a, b) => vals[*a] / vals[*b],
1048 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
1049 TapeOp::Neg(a) => -vals[*a],
1050 TapeOp::Abs(a) => vals[*a].abs(),
1051 TapeOp::Sqrt(a) => vals[*a].sqrt(),
1052 TapeOp::Exp(a) => vals[*a].exp(),
1053 TapeOp::Log(a) => vals[*a].ln(),
1054 TapeOp::Log10(a) => vals[*a].log10(),
1055 TapeOp::Sin(a) => vals[*a].sin(),
1056 TapeOp::Cos(a) => vals[*a].cos(),
1057 TapeOp::Tan(a) => vals[*a].tan(),
1058 TapeOp::Atan(a) => vals[*a].atan(),
1059 TapeOp::Acos(a) => vals[*a].acos(),
1060 TapeOp::Sinh(a) => vals[*a].sinh(),
1061 TapeOp::Cosh(a) => vals[*a].cosh(),
1062 TapeOp::Tanh(a) => vals[*a].tanh(),
1063 TapeOp::Asin(a) => vals[*a].asin(),
1064 TapeOp::Acosh(a) => vals[*a].acosh(),
1065 TapeOp::Asinh(a) => vals[*a].asinh(),
1066 TapeOp::Atanh(a) => atanh(vals[*a]),
1067 TapeOp::Erf(a) => erf(vals[*a]),
1068 TapeOp::XLogX(a) => xlogx(vals[*a]),
1069 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
1070 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
1071 TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
1072 TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
1073 TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
1074 TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
1075 TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
1076 TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
1077 TapeOp::Select(c, t, e) => {
1078 if vals[*c] != 0.0 {
1079 vals[*t]
1080 } else {
1081 vals[*e]
1082 }
1083 }
1084 TapeOp::Funcall(fc) => {
1085 let FuncallData { lib, name, args } = fc.as_ref();
1086 let call_args = funcall_to_ext_args(args, &*vals);
1087 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
1088 res.value
1089 }
1090 };
1091 }
1092 }
1093
1094 pub fn eval_into(&self, x: &[f64], vals: &mut [f64]) -> f64 {
1105 let n = self.ops.len();
1106 if n == 0 {
1107 return 0.0;
1108 }
1109 self.forward_into(x, vals);
1110 vals[n - 1]
1111 }
1112
1113 pub fn hessian_directional(
1130 &self,
1131 vals: &[f64],
1132 seed: &[f64],
1133 weight: f64,
1134 out: &mut [f64],
1135 dot: &mut [f64],
1136 adj: &mut [f64],
1137 adj_dot: &mut [f64],
1138 ) {
1139 let n = self.ops.len();
1140 if n == 0 || weight == 0.0 {
1141 return;
1142 }
1143 debug_assert!(vals.len() >= n);
1144 debug_assert!(dot.len() >= n);
1145 debug_assert!(adj.len() >= n);
1146 debug_assert!(adj_dot.len() >= n);
1147
1148 for i in 0..n {
1152 dot[i] = match &self.ops[i] {
1153 TapeOp::Const(_) => 0.0,
1154 TapeOp::Var(k) => seed[*k],
1155 TapeOp::Add(a, b) => dot[*a] + dot[*b],
1156 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
1157 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
1158 TapeOp::Div(a, b) => {
1159 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
1167 }
1168 TapeOp::Pow(a, b) => {
1169 let u = vals[*a];
1170 let r = vals[*b];
1171 let du = dot[*a];
1172 let dr = dot[*b];
1173 let mut result = 0.0;
1174 if r != 0.0 {
1179 result += r * u.powf(r - 1.0) * du;
1180 }
1181 if u > 0.0 {
1182 result += vals[i] * u.ln() * dr;
1183 }
1184 result
1185 }
1186 TapeOp::Neg(a) => -dot[*a],
1187 TapeOp::Abs(a) => {
1188 if vals[*a] >= 0.0 {
1189 dot[*a]
1190 } else {
1191 -dot[*a]
1192 }
1193 }
1194 TapeOp::Sqrt(a) => {
1195 let sv = vals[i];
1196 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
1197 }
1198 TapeOp::Exp(a) => vals[i] * dot[*a],
1199 TapeOp::Log(a) => dot[*a] / vals[*a],
1200 TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
1201 TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
1202 TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
1203 TapeOp::Tan(a) => {
1204 let t = vals[i];
1205 (1.0 + t * t) * dot[*a]
1206 }
1207 TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
1208 TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
1209 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
1210 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
1211 TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
1212 TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
1213 TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
1214 TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
1215 TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
1216 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
1217 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
1218 TapeOp::CEntropy(a, b) => {
1219 centropy_da(vals[*a], vals[*b]) * dot[*a]
1220 + centropy_db(vals[*a], vals[*b]) * dot[*b]
1221 }
1222 TapeOp::Atan2(a, b) => {
1223 let y = vals[*a];
1224 let x = vals[*b];
1225 let (fa, fb) = atan2_d1(y, x);
1226 fa * dot[*a] + fb * dot[*b]
1227 }
1228 TapeOp::Min(a, b) => {
1230 if vals[*a] <= vals[*b] {
1231 dot[*a]
1232 } else {
1233 dot[*b]
1234 }
1235 }
1236 TapeOp::Max(a, b) => {
1237 if vals[*a] >= vals[*b] {
1238 dot[*a]
1239 } else {
1240 dot[*b]
1241 }
1242 }
1243 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
1244 TapeOp::Select(c, t, e) => {
1245 if vals[*c] != 0.0 {
1246 dot[*t]
1247 } else {
1248 dot[*e]
1249 }
1250 }
1251 TapeOp::Funcall(fc) => {
1252 let FuncallData { lib, name, args } = fc.as_ref();
1253 let call_args = funcall_to_ext_args(args, vals);
1254 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
1255 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1256 let mut acc = 0.0;
1257 let mut k = 0usize;
1258 for arg in args {
1259 if let TapeFuncallArg::Tape(idx) = arg {
1260 acc += derivs[k] * dot[*idx];
1261 k += 1;
1262 }
1263 }
1264 acc
1265 }
1266 };
1267 }
1268
1269 for slot in adj.iter_mut().take(n) {
1273 *slot = 0.0;
1274 }
1275 for slot in adj_dot.iter_mut().take(n) {
1276 *slot = 0.0;
1277 }
1278 adj[n - 1] = 1.0;
1279
1280 for i in (0..n).rev() {
1281 let w = adj[i];
1282 let wd = adj_dot[i];
1283 if w == 0.0 && wd == 0.0 {
1284 continue;
1285 }
1286 match &self.ops[i] {
1287 TapeOp::Const(_) => {}
1288 TapeOp::Var(k) => {
1289 if wd != 0.0 {
1290 out[*k] += weight * wd;
1291 }
1292 }
1293 TapeOp::Add(a, b) => {
1294 adj[*a] += w;
1295 adj[*b] += w;
1296 adj_dot[*a] += wd;
1297 adj_dot[*b] += wd;
1298 }
1299 TapeOp::Sub(a, b) => {
1300 adj[*a] += w;
1301 adj[*b] -= w;
1302 adj_dot[*a] += wd;
1303 adj_dot[*b] -= wd;
1304 }
1305 TapeOp::Mul(a, b) => {
1306 adj[*a] += w * vals[*b];
1307 adj[*b] += w * vals[*a];
1308 adj_dot[*a] += wd * vals[*b] + w * dot[*b];
1309 adj_dot[*b] += wd * vals[*a] + w * dot[*a];
1310 }
1311 TapeOp::Div(a, b) => {
1312 let vb = vals[*b];
1319 let q = vals[i];
1320 let qd = dot[i];
1321 adj[*a] += w / vb;
1322 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1323 adj[*b] -= w * q / vb;
1324 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1325 }
1326 TapeOp::Pow(a, b) => {
1327 let u = vals[*a];
1328 let r = vals[*b];
1329 let du = dot[*a];
1330 let dr = dot[*b];
1331 if r != 0.0 {
1332 if u != 0.0 {
1333 let p_a = r * u.powf(r - 1.0);
1334 adj[*a] += w * p_a;
1335 let mut dp_a = dr * u.powf(r - 1.0);
1336 if u > 0.0 {
1337 dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1338 } else {
1339 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1340 }
1341 adj_dot[*a] += wd * p_a + w * dp_a;
1342 } else if r >= 2.0 {
1343 let p_a = 0.0;
1344 adj[*a] += w * p_a;
1345 let dp_a = if r == 2.0 {
1346 2.0 * du
1347 } else {
1348 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1349 };
1350 adj_dot[*a] += wd * p_a + w * dp_a;
1351 }
1352 }
1353 if u > 0.0 {
1354 let ln_u = u.ln();
1355 let p_b = vals[i] * ln_u;
1356 adj[*b] += w * p_b;
1357 let dur = vals[i] * (r * du / u + dr * ln_u);
1358 let dp_b = dur * ln_u + vals[i] * du / u;
1359 adj_dot[*b] += wd * p_b + w * dp_b;
1360 }
1361 }
1362 TapeOp::Neg(a) => {
1363 adj[*a] -= w;
1364 adj_dot[*a] -= wd;
1365 }
1366 TapeOp::Abs(a) => {
1367 let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
1368 adj[*a] += w * s;
1369 adj_dot[*a] += wd * s;
1370 }
1371 TapeOp::Sqrt(a) => {
1372 let sv = vals[i];
1373 if sv > 0.0 {
1374 let fp = 0.5 / sv;
1375 let fpp = sqrt_d2(vals[*a], sv);
1376 adj[*a] += w * fp;
1377 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1378 }
1379 }
1380 TapeOp::Exp(a) => {
1381 let ev = vals[i];
1382 adj[*a] += w * ev;
1383 adj_dot[*a] += wd * ev + w * ev * dot[*a];
1384 }
1385 TapeOp::Log(a) => {
1386 let u = vals[*a];
1387 adj[*a] += w / u;
1388 adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
1389 }
1390 TapeOp::Log10(a) => {
1391 let u = vals[*a];
1392 adj[*a] += div_u_ln10(w, u);
1393 adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
1394 }
1395 TapeOp::Sin(a) => {
1396 let u = vals[*a];
1397 let cu = u.cos();
1398 adj[*a] += w * cu;
1399 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1400 }
1401 TapeOp::Cos(a) => {
1402 let u = vals[*a];
1403 let su = u.sin();
1404 adj[*a] -= w * su;
1405 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1406 }
1407 TapeOp::Tan(a) => {
1408 let t = vals[i];
1409 let gp = 1.0 + t * t;
1410 let gpp = 2.0 * t * gp;
1411 adj[*a] += w * gp;
1412 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1413 }
1414 TapeOp::Atan(a) => {
1415 let u = vals[*a];
1416 let gp = atan_d1(u);
1417 let gpp = atan_d2(u);
1418 adj[*a] += w * gp;
1419 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1420 }
1421 TapeOp::Acos(a) => {
1422 let u = vals[*a];
1423 let gp = acos_d1(u);
1424 let gpp = acos_d2(u);
1425 adj[*a] += w * gp;
1426 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1427 }
1428 TapeOp::Sinh(a) => {
1429 let u = vals[*a];
1430 let gp = u.cosh();
1431 let gpp = u.sinh();
1432 adj[*a] += w * gp;
1433 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1434 }
1435 TapeOp::Cosh(a) => {
1436 let u = vals[*a];
1437 let gp = u.sinh();
1438 let gpp = u.cosh();
1439 adj[*a] += w * gp;
1440 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1441 }
1442 TapeOp::Tanh(a) => {
1443 let t = vals[i];
1444 let gp = tanh_d1(vals[*a]);
1445 let gpp = -2.0 * t * gp;
1446 adj[*a] += w * gp;
1447 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1448 }
1449 TapeOp::Asin(a) => {
1450 let u = vals[*a];
1451 let gp = asin_d1(u);
1452 let gpp = asin_d2(u);
1453 adj[*a] += w * gp;
1454 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1455 }
1456 TapeOp::Acosh(a) => {
1457 let u = vals[*a];
1458 let gp = acosh_d1(u);
1459 let gpp = acosh_d2(u);
1460 adj[*a] += w * gp;
1461 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1462 }
1463 TapeOp::Asinh(a) => {
1464 let u = vals[*a];
1465 let gp = asinh_d1(u);
1466 let gpp = asinh_d2(u);
1467 adj[*a] += w * gp;
1468 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1469 }
1470 TapeOp::Atanh(a) => {
1471 let u = vals[*a];
1472 let gp = atanh_d1(u);
1473 let gpp = atanh_d2(u);
1474 adj[*a] += w * gp;
1475 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1476 }
1477 TapeOp::Erf(a) => {
1478 let u = vals[*a];
1479 let gp = erf_d1(u);
1480 let gpp = erf_d2(u);
1481 adj[*a] += w * gp;
1482 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1483 }
1484 TapeOp::XLogX(a) => {
1485 let u = vals[*a];
1488 let gp = xlogx_d1(u);
1489 let gpp = xlogx_d2(u);
1490 adj[*a] += w * gp;
1491 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1492 }
1493 TapeOp::CEntropy(a, b) => {
1494 let ua = vals[*a];
1495 let ub = vals[*b];
1496 let fa = centropy_da(ua, ub);
1497 let fb = centropy_db(ua, ub);
1498 let faa = centropy_daa(ua);
1499 let fab = centropy_dab(ub);
1500 let fbb = centropy_dbb(ua, ub);
1501 adj[*a] += w * fa;
1502 adj[*b] += w * fb;
1503 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1504 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1505 }
1506 TapeOp::Atan2(a, b) => {
1507 let y = vals[*a];
1508 let x = vals[*b];
1509 let (fa, fb) = atan2_d1(y, x);
1510 let (faa, fab, fbb) = atan2_d2(y, x);
1511 adj[*a] += w * fa;
1512 adj[*b] += w * fb;
1513 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1514 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1515 }
1516 TapeOp::Min(a, b) => {
1520 let br = if vals[*a] <= vals[*b] { *a } else { *b };
1521 adj[br] += w;
1522 adj_dot[br] += wd;
1523 }
1524 TapeOp::Max(a, b) => {
1525 let br = if vals[*a] >= vals[*b] { *a } else { *b };
1526 adj[br] += w;
1527 adj_dot[br] += wd;
1528 }
1529 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
1531 TapeOp::Select(c, t, e) => {
1534 let br = if vals[*c] != 0.0 { *t } else { *e };
1535 adj[br] += w;
1536 adj_dot[br] += wd;
1537 }
1538 TapeOp::Funcall(fc) => {
1539 let FuncallData { lib, name, args } = fc.as_ref();
1540 let call_args = funcall_to_ext_args(args, vals);
1541 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1542 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1543 let hes = res.hessian.expect("want_hes=true returns hessian");
1544 let real_tape: Vec<usize> = args
1545 .iter()
1546 .filter_map(|a| match a {
1547 TapeFuncallArg::Tape(t) => Some(*t),
1548 TapeFuncallArg::Str(_) => None,
1549 })
1550 .collect();
1551 for (k, &tk) in real_tape.iter().enumerate() {
1552 adj[tk] += w * derivs[k];
1553 let mut second_term = 0.0;
1554 for (l, &tl) in real_tape.iter().enumerate() {
1555 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1556 let h_kl = hes[lo + hi * (hi + 1) / 2];
1557 second_term += h_kl * dot[tl];
1558 }
1559 adj_dot[tk] += wd * derivs[k] + w * second_term;
1560 }
1561 }
1562 }
1563 }
1564 }
1565
1566 pub fn hessian_accumulate(
1573 &self,
1574 x: &[f64],
1575 weight: f64,
1576 hess_map: &HashMap<(usize, usize), usize>,
1577 values: &mut [f64],
1578 ) {
1579 let n = self.ops.len();
1580 if n == 0 || weight == 0.0 {
1581 return;
1582 }
1583 let v = self.forward(x);
1584 let var_indices = self.variables();
1585
1586 let mut dot = vec![0.0f64; n];
1593 let mut adj = vec![0.0f64; n];
1594 let mut adj_dot = vec![0.0f64; n];
1595 for &j in &var_indices {
1596 self.forward_tangent(&v, j, &mut dot);
1597
1598 adj.fill(0.0);
1601 adj_dot.fill(0.0);
1602 adj[n - 1] = 1.0;
1603
1604 for i in (0..n).rev() {
1605 let w = adj[i];
1606 let wd = adj_dot[i];
1607 if w == 0.0 && wd == 0.0 {
1608 continue;
1609 }
1610 match &self.ops[i] {
1611 TapeOp::Const(_) => {}
1612 TapeOp::Var(k) => {
1613 if wd != 0.0 && *k >= j {
1616 if let Some(&pos) = hess_map.get(&(*k, j)) {
1617 values[pos] += weight * wd;
1618 }
1619 }
1620 }
1621 TapeOp::Add(a, b) => {
1622 adj[*a] += w;
1623 adj[*b] += w;
1624 adj_dot[*a] += wd;
1625 adj_dot[*b] += wd;
1626 }
1627 TapeOp::Sub(a, b) => {
1628 adj[*a] += w;
1629 adj[*b] -= w;
1630 adj_dot[*a] += wd;
1631 adj_dot[*b] -= wd;
1632 }
1633 TapeOp::Mul(a, b) => {
1634 adj[*a] += w * v[*b];
1635 adj[*b] += w * v[*a];
1636 adj_dot[*a] += wd * v[*b] + w * dot[*b];
1637 adj_dot[*b] += wd * v[*a] + w * dot[*a];
1638 }
1639 TapeOp::Div(a, b) => {
1640 let vb = v[*b];
1643 let q = v[i];
1644 let qd = dot[i];
1645 adj[*a] += w / vb;
1646 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1647 adj[*b] -= w * q / vb;
1648 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1649 }
1650 TapeOp::Pow(a, b) => {
1651 let u = v[*a];
1652 let r = v[*b];
1653 let du = dot[*a];
1654 let dr = dot[*b];
1655 if r != 0.0 {
1656 if u != 0.0 {
1657 let p_a = r * u.powf(r - 1.0);
1658 adj[*a] += w * p_a;
1659 let mut dp_a = dr * u.powf(r - 1.0);
1660 if u > 0.0 {
1661 dp_a +=
1662 r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1663 } else {
1664 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1665 }
1666 adj_dot[*a] += wd * p_a + w * dp_a;
1667 } else if r >= 2.0 {
1668 let p_a = 0.0;
1669 adj[*a] += w * p_a;
1670 let dp_a = if r == 2.0 {
1671 2.0 * du
1672 } else {
1673 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1674 };
1675 adj_dot[*a] += wd * p_a + w * dp_a;
1676 }
1677 }
1678 if u > 0.0 {
1679 let ln_u = u.ln();
1680 let p_b = v[i] * ln_u;
1681 adj[*b] += w * p_b;
1682 let dur = v[i] * (r * du / u + dr * ln_u);
1683 let dp_b = dur * ln_u + v[i] * du / u;
1684 adj_dot[*b] += wd * p_b + w * dp_b;
1685 }
1686 }
1687 TapeOp::Neg(a) => {
1688 adj[*a] -= w;
1689 adj_dot[*a] -= wd;
1690 }
1691 TapeOp::Abs(a) => {
1692 let s = if v[*a] >= 0.0 { 1.0 } else { -1.0 };
1693 adj[*a] += w * s;
1694 adj_dot[*a] += wd * s;
1695 }
1696 TapeOp::Sqrt(a) => {
1697 let sv = v[i];
1698 if sv > 0.0 {
1699 let fp = 0.5 / sv;
1700 let fpp = sqrt_d2(v[*a], sv);
1701 adj[*a] += w * fp;
1702 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1703 }
1704 }
1705 TapeOp::Exp(a) => {
1706 let ev = v[i];
1707 adj[*a] += w * ev;
1708 adj_dot[*a] += wd * ev + w * ev * dot[*a];
1709 }
1710 TapeOp::Log(a) => {
1711 let u = v[*a];
1712 adj[*a] += w / u;
1713 adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
1714 }
1715 TapeOp::Log10(a) => {
1716 let u = v[*a];
1717 adj[*a] += div_u_ln10(w, u);
1718 adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
1719 }
1720 TapeOp::Sin(a) => {
1721 let u = v[*a];
1722 let cu = u.cos();
1723 adj[*a] += w * cu;
1724 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1725 }
1726 TapeOp::Cos(a) => {
1727 let u = v[*a];
1728 let su = u.sin();
1729 adj[*a] -= w * su;
1730 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1731 }
1732 TapeOp::Tan(a) => {
1733 let t = v[i];
1734 let gp = 1.0 + t * t;
1735 let gpp = 2.0 * t * gp;
1736 adj[*a] += w * gp;
1737 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1738 }
1739 TapeOp::Atan(a) => {
1740 let u = v[*a];
1741 let gp = atan_d1(u);
1742 let gpp = atan_d2(u);
1743 adj[*a] += w * gp;
1744 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1745 }
1746 TapeOp::Acos(a) => {
1747 let u = v[*a];
1748 let gp = acos_d1(u);
1749 let gpp = acos_d2(u);
1750 adj[*a] += w * gp;
1751 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1752 }
1753 TapeOp::Sinh(a) => {
1754 let u = v[*a];
1755 let gp = u.cosh();
1756 let gpp = u.sinh();
1757 adj[*a] += w * gp;
1758 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1759 }
1760 TapeOp::Cosh(a) => {
1761 let u = v[*a];
1762 let gp = u.sinh();
1763 let gpp = u.cosh();
1764 adj[*a] += w * gp;
1765 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1766 }
1767 TapeOp::Tanh(a) => {
1768 let t = v[i];
1769 let gp = tanh_d1(v[*a]);
1770 let gpp = -2.0 * t * gp;
1771 adj[*a] += w * gp;
1772 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1773 }
1774 TapeOp::Asin(a) => {
1775 let u = v[*a];
1776 let gp = asin_d1(u);
1777 let gpp = asin_d2(u);
1778 adj[*a] += w * gp;
1779 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1780 }
1781 TapeOp::Acosh(a) => {
1782 let u = v[*a];
1783 let gp = acosh_d1(u);
1784 let gpp = acosh_d2(u);
1785 adj[*a] += w * gp;
1786 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1787 }
1788 TapeOp::Asinh(a) => {
1789 let u = v[*a];
1790 let gp = asinh_d1(u);
1791 let gpp = asinh_d2(u);
1792 adj[*a] += w * gp;
1793 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1794 }
1795 TapeOp::Atanh(a) => {
1796 let u = v[*a];
1797 let gp = atanh_d1(u);
1798 let gpp = atanh_d2(u);
1799 adj[*a] += w * gp;
1800 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1801 }
1802 TapeOp::Erf(a) => {
1803 let u = v[*a];
1804 let gp = erf_d1(u);
1805 let gpp = erf_d2(u);
1806 adj[*a] += w * gp;
1807 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1808 }
1809 TapeOp::XLogX(a) => {
1810 let u = v[*a];
1813 let gp = xlogx_d1(u);
1814 let gpp = xlogx_d2(u);
1815 adj[*a] += w * gp;
1816 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1817 }
1818 TapeOp::CEntropy(a, b) => {
1819 let ua = v[*a];
1820 let ub = v[*b];
1821 let fa = centropy_da(ua, ub);
1822 let fb = centropy_db(ua, ub);
1823 let faa = centropy_daa(ua);
1824 let fab = centropy_dab(ub);
1825 let fbb = centropy_dbb(ua, ub);
1826 adj[*a] += w * fa;
1827 adj[*b] += w * fb;
1828 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1829 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1830 }
1831 TapeOp::Atan2(a, b) => {
1832 let y = v[*a];
1833 let x = v[*b];
1834 let (fa, fb) = atan2_d1(y, x);
1835 let (faa, fab, fbb) = atan2_d2(y, x);
1836 adj[*a] += w * fa;
1837 adj[*b] += w * fb;
1838 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1839 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1840 }
1841 TapeOp::Min(a, b) => {
1845 let br = if v[*a] <= v[*b] { *a } else { *b };
1846 adj[br] += w;
1847 adj_dot[br] += wd;
1848 }
1849 TapeOp::Max(a, b) => {
1850 let br = if v[*a] >= v[*b] { *a } else { *b };
1851 adj[br] += w;
1852 adj_dot[br] += wd;
1853 }
1854 TapeOp::Cmp(_, _, _)
1856 | TapeOp::And(_, _)
1857 | TapeOp::Or(_, _)
1858 | TapeOp::Not(_) => {}
1859 TapeOp::Select(c, t, e) => {
1862 let br = if v[*c] != 0.0 { *t } else { *e };
1863 adj[br] += w;
1864 adj_dot[br] += wd;
1865 }
1866 TapeOp::Funcall(fc) => {
1867 let FuncallData { lib, name, args } = fc.as_ref();
1868 let call_args = funcall_to_ext_args(args, &v);
1869 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1870 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1871 let hes = res.hessian.expect("want_hes=true returns hessian");
1872 let real_tape: Vec<usize> = args
1873 .iter()
1874 .filter_map(|a| match a {
1875 TapeFuncallArg::Tape(t) => Some(*t),
1876 TapeFuncallArg::Str(_) => None,
1877 })
1878 .collect();
1879 for (k, &tk) in real_tape.iter().enumerate() {
1880 adj[tk] += w * derivs[k];
1881 let mut second_term = 0.0;
1882 for (l, &tl) in real_tape.iter().enumerate() {
1883 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1884 let h_kl = hes[lo + hi * (hi + 1) / 2];
1885 second_term += h_kl * dot[tl];
1886 }
1887 adj_dot[tk] += wd * derivs[k] + w * second_term;
1888 }
1889 }
1890 }
1891 }
1892 }
1893 }
1894
1895 pub fn hessian_sparsity(&self) -> BTreeSet<(usize, usize)> {
1900 let n = self.ops.len();
1901 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
1902 let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
1903
1904 let mut last_use: Vec<usize> = (0..n).collect();
1919 for (i, op) in self.ops.iter().enumerate() {
1920 for_each_input(op, |a| last_use[a] = i);
1921 }
1922 macro_rules! merge {
1926 ($a:expr, $b:expr, $i:expr) => {{
1927 let (a, b, i) = ($a, $b, $i);
1928 if a == b {
1929 if last_use[a] == i {
1930 std::mem::take(&mut var_sets[a])
1931 } else {
1932 var_sets[a].clone()
1933 }
1934 } else {
1935 let take_a = last_use[a] == i
1937 && (last_use[b] != i || var_sets[a].len() >= var_sets[b].len());
1938 if take_a {
1939 let mut s = std::mem::take(&mut var_sets[a]);
1940 s.extend(var_sets[b].iter().copied());
1941 s
1942 } else if last_use[b] == i {
1943 let mut s = std::mem::take(&mut var_sets[b]);
1944 s.extend(var_sets[a].iter().copied());
1945 s
1946 } else {
1947 var_sets[a].union(&var_sets[b]).copied().collect()
1948 }
1949 }
1950 }};
1951 }
1952 macro_rules! carry {
1954 ($a:expr, $i:expr) => {{
1955 let (a, i) = ($a, $i);
1956 if last_use[a] == i {
1957 std::mem::take(&mut var_sets[a])
1958 } else {
1959 var_sets[a].clone()
1960 }
1961 }};
1962 }
1963
1964 let emit_cross =
1965 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1966 for &v1 in s1 {
1967 for &v2 in s2 {
1968 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
1969 pairs.insert((r, c));
1970 }
1971 }
1972 };
1973 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1974 let vars: Vec<usize> = s.iter().copied().collect();
1975 for (ai, &vi) in vars.iter().enumerate() {
1976 for &vj in &vars[..=ai] {
1977 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
1978 pairs.insert((r, c));
1979 }
1980 }
1981 };
1982
1983 for (i, op) in self.ops.iter().enumerate() {
1984 let vset = match op {
1985 TapeOp::Const(_) => BTreeSet::new(),
1986 TapeOp::Var(j) => {
1987 let mut s = BTreeSet::new();
1988 s.insert(*j);
1989 s
1990 }
1991 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => merge!(*a, *b, i),
1992 TapeOp::Neg(a) | TapeOp::Abs(a) => carry!(*a, i),
1993 TapeOp::Mul(a, b) => {
1994 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1995 merge!(*a, *b, i)
1996 }
1997 TapeOp::Div(a, b) => {
1998 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1999 emit_self(&var_sets[*b], &mut pairs);
2000 merge!(*a, *b, i)
2001 }
2002 TapeOp::Pow(a, b) => {
2003 let combined = merge!(*a, *b, i);
2004 emit_self(&combined, &mut pairs);
2005 combined
2006 }
2007 TapeOp::Sqrt(a)
2008 | TapeOp::Exp(a)
2009 | TapeOp::Log(a)
2010 | TapeOp::Log10(a)
2011 | TapeOp::Sin(a)
2012 | TapeOp::Cos(a)
2013 | TapeOp::Tan(a)
2014 | TapeOp::Atan(a)
2015 | TapeOp::Acos(a)
2016 | TapeOp::Sinh(a)
2017 | TapeOp::Cosh(a)
2018 | TapeOp::Tanh(a)
2019 | TapeOp::Asin(a)
2020 | TapeOp::Acosh(a)
2021 | TapeOp::Asinh(a)
2022 | TapeOp::Erf(a)
2023 | TapeOp::XLogX(a)
2024 | TapeOp::Atanh(a) => {
2025 emit_self(&var_sets[*a], &mut pairs);
2026 carry!(*a, i)
2027 }
2028 TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
2033 let combined = merge!(*a, *b, i);
2034 emit_self(&combined, &mut pairs);
2035 combined
2036 }
2037 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
2043 BTreeSet::new()
2044 }
2045 TapeOp::Select(_c, t, e) => merge!(*t, *e, i),
2052 TapeOp::Min(a, b) | TapeOp::Max(a, b) => merge!(*a, *b, i),
2059 TapeOp::Funcall(fc) => {
2060 let args = &fc.args;
2061 let mut combined: BTreeSet<usize> = BTreeSet::new();
2062 for arg in args {
2063 if let TapeFuncallArg::Tape(t) = arg {
2064 for &vv in &var_sets[*t] {
2065 combined.insert(vv);
2066 }
2067 }
2068 }
2069 emit_self(&combined, &mut pairs);
2070 combined
2071 }
2072 };
2073 for_each_input(op, |a| {
2079 if last_use[a] == i {
2080 var_sets[a].clear();
2081 }
2082 });
2083 var_sets.push(vset);
2084 }
2085 pairs
2086 }
2087}
2088
2089fn for_each_input(op: &TapeOp, mut f: impl FnMut(usize)) {
2100 match op {
2101 TapeOp::Const(_) | TapeOp::Var(_) => {}
2102 TapeOp::Neg(a)
2103 | TapeOp::Abs(a)
2104 | TapeOp::Sqrt(a)
2105 | TapeOp::Exp(a)
2106 | TapeOp::Log(a)
2107 | TapeOp::Log10(a)
2108 | TapeOp::Sin(a)
2109 | TapeOp::Cos(a)
2110 | TapeOp::Tan(a)
2111 | TapeOp::Atan(a)
2112 | TapeOp::Acos(a)
2113 | TapeOp::Sinh(a)
2114 | TapeOp::Cosh(a)
2115 | TapeOp::Tanh(a)
2116 | TapeOp::Asin(a)
2117 | TapeOp::Acosh(a)
2118 | TapeOp::Asinh(a)
2119 | TapeOp::Atanh(a)
2120 | TapeOp::Erf(a)
2121 | TapeOp::XLogX(a)
2122 | TapeOp::Not(a) => f(*a),
2123 TapeOp::Add(a, b)
2124 | TapeOp::Sub(a, b)
2125 | TapeOp::Mul(a, b)
2126 | TapeOp::Div(a, b)
2127 | TapeOp::Pow(a, b)
2128 | TapeOp::CEntropy(a, b)
2129 | TapeOp::Atan2(a, b)
2130 | TapeOp::Min(a, b)
2131 | TapeOp::Max(a, b)
2132 | TapeOp::And(a, b)
2133 | TapeOp::Or(a, b)
2134 | TapeOp::Cmp(_, a, b) => {
2135 f(*a);
2136 f(*b);
2137 }
2138 TapeOp::Select(c, t, e) => {
2139 f(*c);
2140 f(*t);
2141 f(*e);
2142 }
2143 TapeOp::Funcall(fc) => {
2144 for arg in &fc.args {
2145 if let TapeFuncallArg::Tape(t) = arg {
2146 f(*t);
2147 }
2148 }
2149 }
2150 }
2151}
2152
2153fn build_recursive(
2154 expr: &Expr,
2155 ops: &mut Vec<TapeOp>,
2156 cache: &mut HashMap<*const Expr, usize>,
2157 resolver: &ExternalResolver,
2158) -> usize {
2159 match expr {
2160 Expr::Const(c) => {
2161 let idx = ops.len();
2162 ops.push(TapeOp::Const(*c));
2163 idx
2164 }
2165 Expr::Var(i) => {
2166 let idx = ops.len();
2167 ops.push(TapeOp::Var(*i));
2168 idx
2169 }
2170 Expr::Binary(op, a, b) => {
2171 if let BinOp::Pow = op {
2179 if let Some(c) = peek_const(b) {
2180 if let Some(idx) = try_emit_const_pow(a, c, ops, cache, resolver) {
2181 return idx;
2182 }
2183 }
2184 }
2185 let l = build_recursive(a, ops, cache, resolver);
2186 let r = build_recursive(b, ops, cache, resolver);
2187 let idx = ops.len();
2188 ops.push(match op {
2189 BinOp::Add => TapeOp::Add(l, r),
2190 BinOp::Sub => TapeOp::Sub(l, r),
2191 BinOp::Mul => TapeOp::Mul(l, r),
2192 BinOp::Div => TapeOp::Div(l, r),
2193 BinOp::Pow => TapeOp::Pow(l, r),
2194 BinOp::Atan2 => TapeOp::Atan2(l, r),
2195 BinOp::CEntropy => TapeOp::CEntropy(l, r),
2196 });
2197 idx
2198 }
2199 Expr::Unary(op, a) => {
2200 let v = build_recursive(a, ops, cache, resolver);
2201 let idx = ops.len();
2202 ops.push(match op {
2203 UnaryOp::Neg => TapeOp::Neg(v),
2204 UnaryOp::Sqrt => TapeOp::Sqrt(v),
2205 UnaryOp::Log => TapeOp::Log(v),
2206 UnaryOp::Log10 => TapeOp::Log10(v),
2207 UnaryOp::Exp => TapeOp::Exp(v),
2208 UnaryOp::Abs => TapeOp::Abs(v),
2209 UnaryOp::Sin => TapeOp::Sin(v),
2210 UnaryOp::Cos => TapeOp::Cos(v),
2211 UnaryOp::Tan => TapeOp::Tan(v),
2212 UnaryOp::Atan => TapeOp::Atan(v),
2213 UnaryOp::Acos => TapeOp::Acos(v),
2214 UnaryOp::Sinh => TapeOp::Sinh(v),
2215 UnaryOp::Cosh => TapeOp::Cosh(v),
2216 UnaryOp::Tanh => TapeOp::Tanh(v),
2217 UnaryOp::Asin => TapeOp::Asin(v),
2218 UnaryOp::Acosh => TapeOp::Acosh(v),
2219 UnaryOp::Asinh => TapeOp::Asinh(v),
2220 UnaryOp::Atanh => TapeOp::Atanh(v),
2221 UnaryOp::Erf => TapeOp::Erf(v),
2222 UnaryOp::XLogX => TapeOp::XLogX(v),
2223 });
2224 idx
2225 }
2226 Expr::Sum(args) => {
2227 if args.is_empty() {
2228 let idx = ops.len();
2229 ops.push(TapeOp::Const(0.0));
2230 return idx;
2231 }
2232 let mut acc = build_recursive(&args[0], ops, cache, resolver);
2233 for a in &args[1..] {
2234 let next = build_recursive(a, ops, cache, resolver);
2235 let idx = ops.len();
2236 ops.push(TapeOp::Add(acc, next));
2237 acc = idx;
2238 }
2239 acc
2240 }
2241 Expr::MinList(args) | Expr::MaxList(args) => {
2249 let is_min = matches!(expr, Expr::MinList(_));
2250 if args.is_empty() {
2251 let idx = ops.len();
2252 ops.push(TapeOp::Const(0.0));
2253 return idx;
2254 }
2255 let mut acc = build_recursive(&args[0], ops, cache, resolver);
2256 for a in &args[1..] {
2257 let next = build_recursive(a, ops, cache, resolver);
2258 let idx = ops.len();
2259 ops.push(if is_min {
2260 TapeOp::Min(acc, next)
2261 } else {
2262 TapeOp::Max(acc, next)
2263 });
2264 acc = idx;
2265 }
2266 acc
2267 }
2268 Expr::Cse(body) => {
2269 let key = Arc::as_ptr(body) as *const Expr;
2276 if let Some(&idx) = cache.get(&key) {
2277 idx
2278 } else {
2279 let idx = build_recursive(body, ops, cache, resolver);
2280 cache.insert(key, idx);
2281 idx
2282 }
2283 }
2284 Expr::Compare(op, a, b) => {
2285 let l = build_recursive(a, ops, cache, resolver);
2286 let r = build_recursive(b, ops, cache, resolver);
2287 let idx = ops.len();
2288 ops.push(TapeOp::Cmp(*op, l, r));
2289 idx
2290 }
2291 Expr::And(a, b) => {
2292 let l = build_recursive(a, ops, cache, resolver);
2293 let r = build_recursive(b, ops, cache, resolver);
2294 let idx = ops.len();
2295 ops.push(TapeOp::And(l, r));
2296 idx
2297 }
2298 Expr::Or(a, b) => {
2299 let l = build_recursive(a, ops, cache, resolver);
2300 let r = build_recursive(b, ops, cache, resolver);
2301 let idx = ops.len();
2302 ops.push(TapeOp::Or(l, r));
2303 idx
2304 }
2305 Expr::Not(a) => {
2306 let v = build_recursive(a, ops, cache, resolver);
2307 let idx = ops.len();
2308 ops.push(TapeOp::Not(v));
2309 idx
2310 }
2311 Expr::Cond { cond, then_, else_ } => {
2312 let c = build_recursive(cond, ops, cache, resolver);
2313 let t = build_recursive(then_, ops, cache, resolver);
2314 let e = build_recursive(else_, ops, cache, resolver);
2315 let idx = ops.len();
2316 ops.push(TapeOp::Select(c, t, e));
2317 idx
2318 }
2319 Expr::Funcall { id, args } => {
2320 let (lib, name) = resolver
2321 .funcs_by_id
2322 .get(id)
2323 .unwrap_or_else(|| panic!("unresolved AMPL funcall id {id}"));
2324 let tape_args: Vec<TapeFuncallArg> = args
2325 .iter()
2326 .map(|a| match a {
2327 FuncallArg::Real(e) => {
2328 TapeFuncallArg::Tape(build_recursive(e, ops, cache, resolver))
2329 }
2330 FuncallArg::Str(s) => TapeFuncallArg::Str(s.clone()),
2331 })
2332 .collect();
2333 let idx = ops.len();
2334 ops.push(TapeOp::Funcall(Box::new(FuncallData {
2335 lib: Arc::clone(lib),
2336 name: name.clone(),
2337 args: tape_args,
2338 })));
2339 idx
2340 }
2341 }
2342}
2343
2344fn peek_const(e: &Expr) -> Option<f64> {
2348 match e {
2349 Expr::Const(c) => Some(*c),
2350 Expr::Cse(body) => peek_const(body),
2351 _ => None,
2352 }
2353}
2354
2355fn try_emit_const_pow(
2363 base_expr: &Expr,
2364 c: f64,
2365 ops: &mut Vec<TapeOp>,
2366 cache: &mut HashMap<*const Expr, usize>,
2367 resolver: &ExternalResolver,
2368) -> Option<usize> {
2369 if c == 0.0 {
2370 let idx = ops.len();
2371 ops.push(TapeOp::Const(1.0));
2372 return Some(idx);
2373 }
2374 if c == 1.0 {
2375 return Some(build_recursive(base_expr, ops, cache, resolver));
2376 }
2377 if c == 0.5 {
2378 let b = build_recursive(base_expr, ops, cache, resolver);
2379 let idx = ops.len();
2380 ops.push(TapeOp::Sqrt(b));
2381 return Some(idx);
2382 }
2383 if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
2388 let n = c.abs() as u32;
2389 if n == 0 {
2390 let idx = ops.len();
2392 ops.push(TapeOp::Const(1.0));
2393 return Some(idx);
2394 }
2395 let b = build_recursive(base_expr, ops, cache, resolver);
2396 let pos = emit_int_pow(b, n, ops);
2397 if c < 0.0 {
2398 let one_idx = ops.len();
2401 ops.push(TapeOp::Const(1.0));
2402 let idx = ops.len();
2403 ops.push(TapeOp::Div(one_idx, pos));
2404 return Some(idx);
2405 }
2406 return Some(pos);
2407 }
2408 None
2409}
2410
2411fn emit_int_pow(base: usize, n: u32, ops: &mut Vec<TapeOp>) -> usize {
2415 debug_assert!(n >= 1);
2416 if n == 1 {
2417 return base;
2418 }
2419 let half = emit_int_pow(base, n / 2, ops);
2420 let squared = ops.len();
2421 ops.push(TapeOp::Mul(half, half));
2422 if n % 2 == 1 {
2423 let idx = ops.len();
2424 ops.push(TapeOp::Mul(squared, base));
2425 idx
2426 } else {
2427 squared
2428 }
2429}
2430
2431#[derive(Debug, Clone)]
2459pub enum SummandOp {
2460 Local(TapeOp),
2463 Shared(usize),
2467}
2468
2469#[derive(Debug, Clone)]
2470pub struct Summand {
2471 pub ops: Vec<SummandOp>,
2472 pub root_slot: usize,
2474 pub local_reach: Vec<usize>,
2476 pub prelude_reach: Vec<usize>,
2479 pub local_vars: Vec<usize>,
2481 pub prelude_vars: Vec<usize>,
2483 pub all_vars: Vec<usize>,
2490}
2491
2492#[derive(Debug, Clone)]
2493pub struct HybridTape {
2494 pub prelude: Vec<TapeOp>,
2499 pub summands: Vec<Summand>,
2500}
2501
2502impl HybridTape {
2503 pub fn build_multi(exprs: &[Expr]) -> Self {
2508 let mut cse_count: HashMap<*const Expr, usize> = HashMap::new();
2512 for e in exprs {
2513 let mut seen_in_root: HashSet<*const Expr> = HashSet::new();
2514 count_cse_appearances(e, &mut seen_in_root, &mut cse_count);
2515 }
2516
2517 let mut prelude: Vec<TapeOp> = Vec::new();
2522 let mut prelude_map: HashMap<*const Expr, usize> = HashMap::new();
2523 let mut summands: Vec<Summand> = Vec::with_capacity(exprs.len());
2524 for e in exprs {
2525 let mut local: Vec<SummandOp> = Vec::new();
2526 let mut local_cache: HashMap<*const Expr, usize> = HashMap::new();
2527 let root_slot = build_into_summand(
2528 e,
2529 &mut local,
2530 &mut local_cache,
2531 &mut prelude,
2532 &mut prelude_map,
2533 &cse_count,
2534 );
2535 summands.push(Summand {
2536 ops: local,
2537 root_slot,
2538 local_reach: Vec::new(),
2539 prelude_reach: Vec::new(),
2540 local_vars: Vec::new(),
2541 prelude_vars: Vec::new(),
2542 all_vars: Vec::new(),
2543 });
2544 }
2545
2546 let mut p_visited: Vec<u32> = vec![0; prelude.len()];
2550 let mut p_epoch: u32 = 0;
2551 let mut p_stack: Vec<usize> = Vec::new();
2552 for s in &mut summands {
2553 let (local_reach, shared_refs) = compute_local_reach(&s.ops, s.root_slot);
2554 s.local_reach = local_reach;
2555
2556 let mut lv: BTreeSet<usize> = BTreeSet::new();
2557 for &i in &s.local_reach {
2558 if let SummandOp::Local(TapeOp::Var(j)) = &s.ops[i] {
2559 lv.insert(*j);
2560 }
2561 }
2562 s.local_vars = lv.iter().copied().collect();
2563
2564 if !shared_refs.is_empty() {
2565 p_epoch += 1;
2566 let mut preach: Vec<usize> = Vec::new();
2567 for &start in &shared_refs {
2568 bfs_prelude(
2569 &prelude,
2570 start,
2571 &mut p_visited,
2572 p_epoch,
2573 &mut p_stack,
2574 &mut preach,
2575 );
2576 }
2577 preach.sort_unstable();
2578 s.prelude_vars = vars_in(&prelude, &preach);
2579 s.prelude_reach = preach;
2580 }
2581
2582 let mut av: BTreeSet<usize> = lv;
2583 for &v in &s.prelude_vars {
2584 av.insert(v);
2585 }
2586 s.all_vars = av.into_iter().collect();
2587 }
2588
2589 HybridTape { prelude, summands }
2590 }
2591
2592 pub fn n_prelude_ops(&self) -> usize {
2593 self.prelude.len()
2594 }
2595 pub fn n_summands(&self) -> usize {
2596 self.summands.len()
2597 }
2598 pub fn max_summand_ops(&self) -> usize {
2599 self.summands.iter().map(|s| s.ops.len()).max().unwrap_or(0)
2600 }
2601 pub fn total_local_ops(&self) -> usize {
2602 self.summands.iter().map(|s| s.ops.len()).sum()
2603 }
2604
2605 pub fn forward_prelude(&self, x: &[f64], prelude_vals: &mut [f64]) {
2608 debug_assert_eq!(prelude_vals.len(), self.prelude.len());
2609 for i in 0..self.prelude.len() {
2610 prelude_vals[i] = fwd_step(&self.prelude[i], x, prelude_vals);
2611 }
2612 }
2613
2614 pub fn forward_summand(
2617 &self,
2618 s: &Summand,
2619 x: &[f64],
2620 prelude_vals: &[f64],
2621 local_vals: &mut [f64],
2622 ) {
2623 debug_assert!(local_vals.len() >= s.ops.len());
2624 for i in 0..s.ops.len() {
2625 local_vals[i] = match &s.ops[i] {
2626 SummandOp::Local(op) => fwd_step(op, x, local_vals),
2627 SummandOp::Shared(k) => prelude_vals[*k],
2628 };
2629 }
2630 }
2631
2632 #[inline]
2634 pub fn root_value(&self, s: &Summand, local_vals: &[f64]) -> f64 {
2635 local_vals[s.root_slot]
2636 }
2637
2638 #[allow(clippy::too_many_arguments)]
2645 pub fn gradient_summand(
2646 &self,
2647 s: &Summand,
2648 prelude_vals: &[f64],
2649 local_vals: &[f64],
2650 seed: f64,
2651 grad: &mut [f64],
2652 local_adj: &mut [f64],
2653 prelude_adj: &mut [f64],
2654 ) {
2655 if seed == 0.0 || s.local_reach.is_empty() {
2656 return;
2657 }
2658 for &i in &s.local_reach {
2659 local_adj[i] = 0.0;
2660 }
2661 for &i in &s.prelude_reach {
2662 prelude_adj[i] = 0.0;
2663 }
2664 local_adj[s.root_slot] = seed;
2665 for &i in s.local_reach.iter().rev() {
2666 let a = local_adj[i];
2667 if a == 0.0 {
2668 continue;
2669 }
2670 match &s.ops[i] {
2671 SummandOp::Local(op) => rev_step(op, i, local_vals, local_adj, a, grad),
2672 SummandOp::Shared(k) => {
2673 prelude_adj[*k] += a;
2674 }
2675 }
2676 }
2677 for &i in s.prelude_reach.iter().rev() {
2678 let a = prelude_adj[i];
2679 if a == 0.0 {
2680 continue;
2681 }
2682 rev_step(&self.prelude[i], i, prelude_vals, prelude_adj, a, grad);
2683 }
2684 }
2685
2686 pub fn prelude_tangent(
2706 &self,
2707 prelude_vals: &[f64],
2708 seed: &[f64],
2709 reach: &[u32],
2710 prelude_dot: &mut [f64],
2711 ) {
2712 debug_assert!(prelude_dot.len() >= self.prelude.len());
2713 for &i in reach {
2714 let i = i as usize;
2715 prelude_dot[i] = fwd_dir_step(&self.prelude[i], seed, prelude_vals, prelude_dot, i);
2716 }
2717 }
2718
2719 #[allow(clippy::too_many_arguments)]
2747 pub fn hessian_summand_directional(
2748 &self,
2749 s: &Summand,
2750 local_vals: &[f64],
2751 prelude_dot: &[f64],
2752 seed: &[f64],
2753 weight: f64,
2754 out: &mut [f64],
2755 local_dot: &mut [f64],
2756 local_adj: &mut [f64],
2757 local_adj_dot: &mut [f64],
2758 prelude_adj: &mut [f64],
2759 prelude_adj_dot: &mut [f64],
2760 ) {
2761 if weight == 0.0 || s.local_reach.is_empty() {
2762 return;
2763 }
2764 for &i in &s.local_reach {
2765 local_adj[i] = 0.0;
2766 local_adj_dot[i] = 0.0;
2767 }
2768 for &i in &s.local_reach {
2772 local_dot[i] = match &s.ops[i] {
2773 SummandOp::Local(op) => fwd_dir_step(op, seed, local_vals, local_dot, i),
2774 SummandOp::Shared(k) => prelude_dot[*k],
2775 };
2776 }
2777 local_adj[s.root_slot] = 1.0;
2778 for &i in s.local_reach.iter().rev() {
2779 let w = local_adj[i];
2780 let wd = local_adj_dot[i];
2781 if w == 0.0 && wd == 0.0 {
2782 continue;
2783 }
2784 match &s.ops[i] {
2785 SummandOp::Local(op) => {
2786 ror_dir_step(
2787 op,
2788 i,
2789 local_vals,
2790 local_dot,
2791 local_adj,
2792 local_adj_dot,
2793 w,
2794 wd,
2795 weight,
2796 out,
2797 );
2798 }
2799 SummandOp::Shared(k) => {
2800 prelude_adj[*k] += weight * w;
2801 prelude_adj_dot[*k] += weight * wd;
2802 }
2803 }
2804 }
2805 }
2806
2807 pub fn prelude_reverse_directional(
2821 &self,
2822 prelude_vals: &[f64],
2823 prelude_dot: &[f64],
2824 reach: &[u32],
2825 out: &mut [f64],
2826 prelude_adj: &mut [f64],
2827 prelude_adj_dot: &mut [f64],
2828 ) {
2829 for &i in reach.iter().rev() {
2830 let i = i as usize;
2831 let w = prelude_adj[i];
2832 let wd = prelude_adj_dot[i];
2833 if w == 0.0 && wd == 0.0 {
2834 continue;
2835 }
2836 prelude_adj[i] = 0.0;
2837 prelude_adj_dot[i] = 0.0;
2838 ror_dir_step(
2839 &self.prelude[i],
2840 i,
2841 prelude_vals,
2842 prelude_dot,
2843 prelude_adj,
2844 prelude_adj_dot,
2845 w,
2846 wd,
2847 1.0,
2848 out,
2849 );
2850 }
2851 }
2852
2853 pub fn hessian_sparsity_all(&self) -> BTreeSet<(usize, usize)> {
2856 let mut pairs = hessian_sparsity_impl(&self.prelude);
2857
2858 let prelude_var_sets = compute_var_sets(&self.prelude);
2861
2862 for s in &self.summands {
2863 summand_sparsity(&s.ops, &prelude_var_sets, &mut pairs);
2864 }
2865 pairs
2866 }
2867}
2868
2869pub fn hybrid_supported(exprs: &[Expr]) -> bool {
2880 let mut stack: Vec<&Expr> = exprs.iter().collect();
2881 let mut seen_cse: HashSet<*const Expr> = HashSet::new();
2882 while let Some(e) = stack.pop() {
2883 match e {
2884 Expr::Const(_) | Expr::Var(_) => {}
2885 Expr::Binary(_, a, b) => {
2886 stack.push(a);
2887 stack.push(b);
2888 }
2889 Expr::Unary(_, a) => stack.push(a),
2890 Expr::Sum(args) => stack.extend(args.iter()),
2891 Expr::Cse(body) => {
2892 if seen_cse.insert(Arc::as_ptr(body)) {
2893 stack.push(body);
2894 }
2895 }
2896 Expr::Compare(..)
2897 | Expr::And(..)
2898 | Expr::Or(..)
2899 | Expr::Not(_)
2900 | Expr::Cond { .. }
2901 | Expr::MinList(_)
2902 | Expr::MaxList(_)
2903 | Expr::Funcall { .. } => return false,
2904 }
2905 }
2906 true
2907}
2908
2909fn cse_contains_funcall(expr: &Expr, built: &HashMap<*const Expr, usize>) -> bool {
2930 let mut stack: Vec<&Expr> = vec![expr];
2931 let mut seen: HashSet<*const Expr> = HashSet::new();
2932 while let Some(e) = stack.pop() {
2933 match e {
2934 Expr::Funcall { .. } => return true,
2935 Expr::Const(_) | Expr::Var(_) => {}
2936 Expr::Binary(_, a, b) | Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2937 stack.push(a);
2938 stack.push(b);
2939 }
2940 Expr::Unary(_, a) | Expr::Not(a) => stack.push(a),
2941 Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2942 stack.extend(args.iter())
2943 }
2944 Expr::Cond { cond, then_, else_ } => {
2945 stack.push(cond);
2946 stack.push(then_);
2947 stack.push(else_);
2948 }
2949 Expr::Cse(body) => {
2950 let key = Arc::as_ptr(body);
2951 if !built.contains_key(&key) && seen.insert(key) {
2952 stack.push(body);
2953 }
2954 }
2955 }
2956 }
2957 false
2958}
2959
2960fn count_cse_appearances(
2965 e: &Expr,
2966 seen_in_root: &mut HashSet<*const Expr>,
2967 counts: &mut HashMap<*const Expr, usize>,
2968) {
2969 match e {
2970 Expr::Const(_) | Expr::Var(_) => {}
2971 Expr::Binary(_, a, b) => {
2972 count_cse_appearances(a, seen_in_root, counts);
2973 count_cse_appearances(b, seen_in_root, counts);
2974 }
2975 Expr::Unary(_, a) => count_cse_appearances(a, seen_in_root, counts),
2976 Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2977 for a in args {
2978 count_cse_appearances(a, seen_in_root, counts);
2979 }
2980 }
2981 Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2982 count_cse_appearances(a, seen_in_root, counts);
2983 count_cse_appearances(b, seen_in_root, counts);
2984 }
2985 Expr::Not(a) => count_cse_appearances(a, seen_in_root, counts),
2986 Expr::Cond { cond, then_, else_ } => {
2987 count_cse_appearances(cond, seen_in_root, counts);
2988 count_cse_appearances(then_, seen_in_root, counts);
2989 count_cse_appearances(else_, seen_in_root, counts);
2990 }
2991 Expr::Cse(body) => {
2992 let key = Arc::as_ptr(body) as *const Expr;
2993 if seen_in_root.insert(key) {
2994 *counts.entry(key).or_insert(0) += 1;
2995 count_cse_appearances(body, seen_in_root, counts);
2996 }
2997 }
2998 Expr::Funcall { args, .. } => {
2999 for arg in args {
3000 if let FuncallArg::Real(e) = arg {
3001 count_cse_appearances(e, seen_in_root, counts);
3002 }
3003 }
3004 }
3005 }
3006}
3007
3008fn build_into_summand(
3014 expr: &Expr,
3015 local: &mut Vec<SummandOp>,
3016 local_cache: &mut HashMap<*const Expr, usize>,
3017 prelude: &mut Vec<TapeOp>,
3018 prelude_map: &mut HashMap<*const Expr, usize>,
3019 cse_count: &HashMap<*const Expr, usize>,
3020) -> usize {
3021 match expr {
3022 Expr::Const(c) => {
3023 let i = local.len();
3024 local.push(SummandOp::Local(TapeOp::Const(*c)));
3025 i
3026 }
3027 Expr::Var(j) => {
3028 let i = local.len();
3029 local.push(SummandOp::Local(TapeOp::Var(*j)));
3030 i
3031 }
3032 Expr::Binary(op, a, b) => {
3033 if let BinOp::Pow = op {
3034 if let Some(c) = peek_const(b) {
3035 if let Some(i) = try_emit_const_pow_summand(
3036 a,
3037 c,
3038 local,
3039 local_cache,
3040 prelude,
3041 prelude_map,
3042 cse_count,
3043 ) {
3044 return i;
3045 }
3046 }
3047 }
3048 let l = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3049 let r = build_into_summand(b, local, local_cache, prelude, prelude_map, cse_count);
3050 let i = local.len();
3051 local.push(SummandOp::Local(match op {
3052 BinOp::Add => TapeOp::Add(l, r),
3053 BinOp::Sub => TapeOp::Sub(l, r),
3054 BinOp::Mul => TapeOp::Mul(l, r),
3055 BinOp::Div => TapeOp::Div(l, r),
3056 BinOp::Pow => TapeOp::Pow(l, r),
3057 BinOp::Atan2 => TapeOp::Atan2(l, r),
3058 BinOp::CEntropy => TapeOp::CEntropy(l, r),
3059 }));
3060 i
3061 }
3062 Expr::Unary(op, a) => {
3063 let v = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3064 let i = local.len();
3065 local.push(SummandOp::Local(match op {
3066 UnaryOp::Neg => TapeOp::Neg(v),
3067 UnaryOp::Sqrt => TapeOp::Sqrt(v),
3068 UnaryOp::Log => TapeOp::Log(v),
3069 UnaryOp::Log10 => TapeOp::Log10(v),
3070 UnaryOp::Exp => TapeOp::Exp(v),
3071 UnaryOp::Abs => TapeOp::Abs(v),
3072 UnaryOp::Sin => TapeOp::Sin(v),
3073 UnaryOp::Cos => TapeOp::Cos(v),
3074 UnaryOp::Tan => TapeOp::Tan(v),
3075 UnaryOp::Atan => TapeOp::Atan(v),
3076 UnaryOp::Acos => TapeOp::Acos(v),
3077 UnaryOp::Sinh => TapeOp::Sinh(v),
3078 UnaryOp::Cosh => TapeOp::Cosh(v),
3079 UnaryOp::Tanh => TapeOp::Tanh(v),
3080 UnaryOp::Asin => TapeOp::Asin(v),
3081 UnaryOp::Acosh => TapeOp::Acosh(v),
3082 UnaryOp::Asinh => TapeOp::Asinh(v),
3083 UnaryOp::Atanh => TapeOp::Atanh(v),
3084 UnaryOp::Erf => TapeOp::Erf(v),
3085 UnaryOp::XLogX => TapeOp::XLogX(v),
3086 }));
3087 i
3088 }
3089 Expr::Sum(args) => {
3090 if args.is_empty() {
3091 let i = local.len();
3092 local.push(SummandOp::Local(TapeOp::Const(0.0)));
3093 return i;
3094 }
3095 let mut acc = build_into_summand(
3096 &args[0],
3097 local,
3098 local_cache,
3099 prelude,
3100 prelude_map,
3101 cse_count,
3102 );
3103 for a in &args[1..] {
3104 let nxt =
3105 build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
3106 let i = local.len();
3107 local.push(SummandOp::Local(TapeOp::Add(acc, nxt)));
3108 acc = i;
3109 }
3110 acc
3111 }
3112 Expr::Cse(body) => {
3113 let key = Arc::as_ptr(body) as *const Expr;
3114 if let Some(&li) = local_cache.get(&key) {
3115 return li;
3116 }
3117 let promoted = cse_count.get(&key).copied().unwrap_or(0) >= 2;
3118 if promoted {
3119 if cse_contains_funcall(body, prelude_map) {
3126 panic!(
3127 "HybridTape: AMPL external function calls are not supported on the \
3128 hybrid (partial-separability) tape path. Build with \
3129 Tape::build_with_externals instead."
3130 );
3131 }
3132 let pslot =
3137 build_recursive(expr, prelude, prelude_map, &ExternalResolver::default());
3138 let li = local.len();
3139 local.push(SummandOp::Shared(pslot));
3140 local_cache.insert(key, li);
3141 li
3142 } else {
3143 let li =
3144 build_into_summand(body, local, local_cache, prelude, prelude_map, cse_count);
3145 local_cache.insert(key, li);
3146 li
3147 }
3148 }
3149 Expr::Compare(_, _, _)
3150 | Expr::And(_, _)
3151 | Expr::Or(_, _)
3152 | Expr::Not(_)
3153 | Expr::Cond { .. }
3154 | Expr::MinList(_)
3155 | Expr::MaxList(_) => {
3156 panic!(
3157 "HybridTape: conditional / logical / min-max opcodes (comparisons, \
3158 AND/OR/NOT, if-then-else, min/max lists) are not supported on the \
3159 hybrid (partial-separability) tape path. Build with \
3160 Tape::build_with_externals instead."
3161 );
3162 }
3163 Expr::Funcall { .. } => {
3164 panic!(
3165 "HybridTape: AMPL external function calls are not supported on the \
3166 hybrid (partial-separability) tape path. Build with Tape::build_with_externals \
3167 instead."
3168 );
3169 }
3170 }
3171}
3172
3173fn try_emit_const_pow_summand(
3176 base_expr: &Expr,
3177 c: f64,
3178 local: &mut Vec<SummandOp>,
3179 local_cache: &mut HashMap<*const Expr, usize>,
3180 prelude: &mut Vec<TapeOp>,
3181 prelude_map: &mut HashMap<*const Expr, usize>,
3182 cse_count: &HashMap<*const Expr, usize>,
3183) -> Option<usize> {
3184 if c == 0.0 {
3185 let i = local.len();
3186 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3187 return Some(i);
3188 }
3189 if c == 1.0 {
3190 return Some(build_into_summand(
3191 base_expr,
3192 local,
3193 local_cache,
3194 prelude,
3195 prelude_map,
3196 cse_count,
3197 ));
3198 }
3199 if c == 0.5 {
3200 let b = build_into_summand(
3201 base_expr,
3202 local,
3203 local_cache,
3204 prelude,
3205 prelude_map,
3206 cse_count,
3207 );
3208 let i = local.len();
3209 local.push(SummandOp::Local(TapeOp::Sqrt(b)));
3210 return Some(i);
3211 }
3212 if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
3213 let n = c.abs() as u32;
3214 if n == 0 {
3215 let i = local.len();
3216 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3217 return Some(i);
3218 }
3219 let b = build_into_summand(
3220 base_expr,
3221 local,
3222 local_cache,
3223 prelude,
3224 prelude_map,
3225 cse_count,
3226 );
3227 let pos = emit_int_pow_summand(b, n, local);
3228 if c < 0.0 {
3229 let one_idx = local.len();
3230 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3231 let i = local.len();
3232 local.push(SummandOp::Local(TapeOp::Div(one_idx, pos)));
3233 return Some(i);
3234 }
3235 return Some(pos);
3236 }
3237 None
3238}
3239
3240fn emit_int_pow_summand(base: usize, n: u32, local: &mut Vec<SummandOp>) -> usize {
3241 debug_assert!(n >= 1);
3242 if n == 1 {
3243 return base;
3244 }
3245 let half = emit_int_pow_summand(base, n / 2, local);
3246 let squared = local.len();
3247 local.push(SummandOp::Local(TapeOp::Mul(half, half)));
3248 if n % 2 == 1 {
3249 let i = local.len();
3250 local.push(SummandOp::Local(TapeOp::Mul(squared, base)));
3251 i
3252 } else {
3253 squared
3254 }
3255}
3256
3257fn compute_local_reach(ops: &[SummandOp], root: usize) -> (Vec<usize>, Vec<usize>) {
3261 let mut visited = vec![false; ops.len()];
3262 let mut reach: Vec<usize> = Vec::new();
3263 let mut shared: BTreeSet<usize> = BTreeSet::new();
3264 let mut stack: Vec<usize> = Vec::with_capacity(16);
3265 visited[root] = true;
3266 reach.push(root);
3267 stack.push(root);
3268 while let Some(s) = stack.pop() {
3269 match &ops[s] {
3270 SummandOp::Local(op) => {
3271 let (a, b) = op_operands(op);
3272 if let Some(a) = a {
3273 if !visited[a] {
3274 visited[a] = true;
3275 reach.push(a);
3276 stack.push(a);
3277 }
3278 }
3279 if let Some(b) = b {
3280 if !visited[b] {
3281 visited[b] = true;
3282 reach.push(b);
3283 stack.push(b);
3284 }
3285 }
3286 }
3287 SummandOp::Shared(k) => {
3288 shared.insert(*k);
3289 }
3290 }
3291 }
3292 reach.sort_unstable();
3293 (reach, shared.into_iter().collect())
3294}
3295
3296fn bfs_prelude(
3300 prelude: &[TapeOp],
3301 start: usize,
3302 visited: &mut [u32],
3303 cur: u32,
3304 stack: &mut Vec<usize>,
3305 out: &mut Vec<usize>,
3306) {
3307 if visited[start] == cur {
3308 return;
3309 }
3310 visited[start] = cur;
3311 out.push(start);
3312 stack.push(start);
3313 while let Some(s) = stack.pop() {
3314 let (a, b) = op_operands(&prelude[s]);
3315 if let Some(a) = a {
3316 if visited[a] != cur {
3317 visited[a] = cur;
3318 out.push(a);
3319 stack.push(a);
3320 }
3321 }
3322 if let Some(b) = b {
3323 if visited[b] != cur {
3324 visited[b] = cur;
3325 out.push(b);
3326 stack.push(b);
3327 }
3328 }
3329 }
3330}
3331
3332fn compute_var_sets(ops: &[TapeOp]) -> Vec<BTreeSet<usize>> {
3336 let mut out: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3337 for op in ops {
3338 let vs: BTreeSet<usize> = match op {
3339 TapeOp::Const(_) => BTreeSet::new(),
3340 TapeOp::Var(j) => {
3341 let mut s = BTreeSet::new();
3342 s.insert(*j);
3343 s
3344 }
3345 TapeOp::Add(a, b)
3346 | TapeOp::Sub(a, b)
3347 | TapeOp::Mul(a, b)
3348 | TapeOp::Div(a, b)
3349 | TapeOp::Pow(a, b)
3350 | TapeOp::Atan2(a, b)
3351 | TapeOp::CEntropy(a, b) => out[*a].union(&out[*b]).copied().collect(),
3352 TapeOp::Neg(a)
3353 | TapeOp::Abs(a)
3354 | TapeOp::Sqrt(a)
3355 | TapeOp::Exp(a)
3356 | TapeOp::Log(a)
3357 | TapeOp::Log10(a)
3358 | TapeOp::Sin(a)
3359 | TapeOp::Cos(a)
3360 | TapeOp::Tan(a)
3361 | TapeOp::Atan(a)
3362 | TapeOp::Acos(a)
3363 | TapeOp::Sinh(a)
3364 | TapeOp::Cosh(a)
3365 | TapeOp::Tanh(a)
3366 | TapeOp::Asin(a)
3367 | TapeOp::Acosh(a)
3368 | TapeOp::Asinh(a)
3369 | TapeOp::Erf(a)
3370 | TapeOp::XLogX(a)
3371 | TapeOp::Atanh(a) => out[*a].clone(),
3372 TapeOp::Cmp(_, _, _)
3373 | TapeOp::And(_, _)
3374 | TapeOp::Or(_, _)
3375 | TapeOp::Not(_)
3376 | TapeOp::Select(_, _, _)
3377 | TapeOp::Min(_, _)
3378 | TapeOp::Max(_, _) => unreachable!(
3379 "HybridTape prelude cannot contain conditional / logical / min-max \
3380 TapeOps; build_into_summand panics on those Expr variants."
3381 ),
3382 TapeOp::Funcall(_) => unreachable!(
3383 "HybridTape prelude cannot contain TapeOp::Funcall; \
3384 build_into_summand panics on Expr::Funcall."
3385 ),
3386 };
3387 out.push(vs);
3388 }
3389 out
3390}
3391
3392fn summand_sparsity(
3397 ops: &[SummandOp],
3398 prelude_var_sets: &[BTreeSet<usize>],
3399 pairs: &mut BTreeSet<(usize, usize)>,
3400) {
3401 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3402 let emit_cross =
3403 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3404 for &v1 in s1 {
3405 for &v2 in s2 {
3406 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
3407 pairs.insert((r, c));
3408 }
3409 }
3410 };
3411 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3412 let vars: Vec<usize> = s.iter().copied().collect();
3413 for (ai, &vi) in vars.iter().enumerate() {
3414 for &vj in &vars[..=ai] {
3415 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
3416 pairs.insert((r, c));
3417 }
3418 }
3419 };
3420 for so in ops {
3421 let vset: BTreeSet<usize> = match so {
3422 SummandOp::Shared(k) => prelude_var_sets[*k].clone(),
3423 SummandOp::Local(op) => match op {
3424 TapeOp::Const(_) => BTreeSet::new(),
3425 TapeOp::Var(j) => {
3426 let mut s = BTreeSet::new();
3427 s.insert(*j);
3428 s
3429 }
3430 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
3431 var_sets[*a].union(&var_sets[*b]).copied().collect()
3432 }
3433 TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
3434 TapeOp::Mul(a, b) => {
3435 emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3436 var_sets[*a].union(&var_sets[*b]).copied().collect()
3437 }
3438 TapeOp::Div(a, b) => {
3439 emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3440 emit_self(&var_sets[*b], pairs);
3441 var_sets[*a].union(&var_sets[*b]).copied().collect()
3442 }
3443 TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
3444 let combined: BTreeSet<usize> =
3445 var_sets[*a].union(&var_sets[*b]).copied().collect();
3446 emit_self(&combined, pairs);
3447 combined
3448 }
3449 TapeOp::Sqrt(a)
3450 | TapeOp::Exp(a)
3451 | TapeOp::Log(a)
3452 | TapeOp::Log10(a)
3453 | TapeOp::Sin(a)
3454 | TapeOp::Cos(a)
3455 | TapeOp::Tan(a)
3456 | TapeOp::Atan(a)
3457 | TapeOp::Acos(a)
3458 | TapeOp::Sinh(a)
3459 | TapeOp::Cosh(a)
3460 | TapeOp::Tanh(a)
3461 | TapeOp::Asin(a)
3462 | TapeOp::Acosh(a)
3463 | TapeOp::Asinh(a)
3464 | TapeOp::Erf(a)
3465 | TapeOp::XLogX(a)
3466 | TapeOp::Atanh(a) => {
3467 emit_self(&var_sets[*a], pairs);
3468 var_sets[*a].clone()
3469 }
3470 TapeOp::Cmp(_, _, _)
3471 | TapeOp::And(_, _)
3472 | TapeOp::Or(_, _)
3473 | TapeOp::Not(_)
3474 | TapeOp::Select(_, _, _)
3475 | TapeOp::Min(_, _)
3476 | TapeOp::Max(_, _) => unreachable!(
3477 "HybridTape summand cannot contain conditional / logical / min-max \
3478 TapeOps; build_into_summand panics on those Expr variants."
3479 ),
3480 TapeOp::Funcall(_) => unreachable!(
3481 "HybridTape summand cannot contain TapeOp::Funcall; \
3482 build_into_summand panics on Expr::Funcall."
3483 ),
3484 },
3485 };
3486 var_sets.push(vset);
3487 }
3488}
3489
3490#[inline]
3493pub(crate) fn op_operands(op: &TapeOp) -> (Option<usize>, Option<usize>) {
3494 match op {
3495 TapeOp::Const(_) | TapeOp::Var(_) => (None, None),
3496 TapeOp::Add(a, b)
3497 | TapeOp::Sub(a, b)
3498 | TapeOp::Mul(a, b)
3499 | TapeOp::Div(a, b)
3500 | TapeOp::Pow(a, b)
3501 | TapeOp::Atan2(a, b)
3502 | TapeOp::CEntropy(a, b) => (Some(*a), Some(*b)),
3503 TapeOp::Neg(a)
3504 | TapeOp::Abs(a)
3505 | TapeOp::Sqrt(a)
3506 | TapeOp::Exp(a)
3507 | TapeOp::Log(a)
3508 | TapeOp::Log10(a)
3509 | TapeOp::Sin(a)
3510 | TapeOp::Cos(a)
3511 | TapeOp::Tan(a)
3512 | TapeOp::Atan(a)
3513 | TapeOp::Acos(a)
3514 | TapeOp::Sinh(a)
3515 | TapeOp::Cosh(a)
3516 | TapeOp::Tanh(a)
3517 | TapeOp::Asin(a)
3518 | TapeOp::Acosh(a)
3519 | TapeOp::Asinh(a)
3520 | TapeOp::Erf(a)
3521 | TapeOp::XLogX(a)
3522 | TapeOp::Atanh(a) => (Some(*a), None),
3523 TapeOp::Cmp(_, a, b) | TapeOp::And(a, b) | TapeOp::Or(a, b) => (Some(*a), Some(*b)),
3529 TapeOp::Not(a) => (Some(*a), None),
3530 TapeOp::Select(_, _, _) => unreachable!(
3531 "op_operands: TapeOp::Select has three operands and is unsupported on \
3532 the HybridTape path"
3533 ),
3534 TapeOp::Min(_, _) | TapeOp::Max(_, _) => unreachable!(
3535 "op_operands: TapeOp::Min/Max are unsupported on the HybridTape path \
3536 (build_into_summand rejects min/max lists)"
3537 ),
3538 TapeOp::Funcall(_) => unreachable!(
3546 "op_operands: TapeOp::Funcall is unsupported on the HybridTape path \
3547 (build_into_summand rejects external function calls)"
3548 ),
3549 }
3550}
3551
3552fn vars_in(ops: &[TapeOp], reach: &[usize]) -> Vec<usize> {
3553 let mut s: BTreeSet<usize> = BTreeSet::new();
3554 for &i in reach {
3555 if let TapeOp::Var(j) = &ops[i] {
3556 s.insert(*j);
3557 }
3558 }
3559 s.into_iter().collect()
3560}
3561
3562#[inline]
3565fn fwd_step(op: &TapeOp, x: &[f64], vals: &[f64]) -> f64 {
3566 match op {
3567 TapeOp::Const(c) => *c,
3568 TapeOp::Var(i) => x[*i],
3569 TapeOp::Add(a, b) => vals[*a] + vals[*b],
3570 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
3571 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
3572 TapeOp::Div(a, b) => vals[*a] / vals[*b],
3573 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
3574 TapeOp::Neg(a) => -vals[*a],
3575 TapeOp::Abs(a) => vals[*a].abs(),
3576 TapeOp::Sqrt(a) => vals[*a].sqrt(),
3577 TapeOp::Exp(a) => vals[*a].exp(),
3578 TapeOp::Log(a) => vals[*a].ln(),
3579 TapeOp::Log10(a) => vals[*a].log10(),
3580 TapeOp::Sin(a) => vals[*a].sin(),
3581 TapeOp::Cos(a) => vals[*a].cos(),
3582 TapeOp::Tan(a) => vals[*a].tan(),
3583 TapeOp::Atan(a) => vals[*a].atan(),
3584 TapeOp::Acos(a) => vals[*a].acos(),
3585 TapeOp::Sinh(a) => vals[*a].sinh(),
3586 TapeOp::Cosh(a) => vals[*a].cosh(),
3587 TapeOp::Tanh(a) => vals[*a].tanh(),
3588 TapeOp::Asin(a) => vals[*a].asin(),
3589 TapeOp::Acosh(a) => vals[*a].acosh(),
3590 TapeOp::Asinh(a) => vals[*a].asinh(),
3591 TapeOp::Atanh(a) => atanh(vals[*a]),
3592 TapeOp::Erf(a) => erf(vals[*a]),
3593 TapeOp::XLogX(a) => xlogx(vals[*a]),
3594 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
3595 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
3596 TapeOp::Cmp(_, _, _)
3597 | TapeOp::And(_, _)
3598 | TapeOp::Or(_, _)
3599 | TapeOp::Not(_)
3600 | TapeOp::Select(_, _, _)
3601 | TapeOp::Min(_, _)
3602 | TapeOp::Max(_, _) => panic!(
3603 "GlobalTape free-function kernels do not implement conditional / logical \
3604 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3605 instead."
3606 ),
3607 TapeOp::Funcall(fc) => {
3608 let FuncallData { lib, name, args } = fc.as_ref();
3609 let call_args = funcall_to_ext_args(args, vals);
3610 let res = lib
3611 .eval(name, &call_args, false, false)
3612 .unwrap_or_else(|e| panic!("external function '{name}' eval failed: {e}"));
3613 res.value
3614 }
3615 }
3616}
3617
3618#[inline]
3619fn rev_step(op: &TapeOp, i: usize, vals: &[f64], adj: &mut [f64], a: f64, grad: &mut [f64]) {
3620 match op {
3621 TapeOp::Const(_) => {}
3622 TapeOp::Var(j) => {
3623 grad[*j] += a;
3624 }
3625 TapeOp::Add(l, r) => {
3626 adj[*l] += a;
3627 adj[*r] += a;
3628 }
3629 TapeOp::Sub(l, r) => {
3630 adj[*l] += a;
3631 adj[*r] -= a;
3632 }
3633 TapeOp::Mul(l, r) => {
3634 adj[*l] += a * vals[*r];
3635 adj[*r] += a * vals[*l];
3636 }
3637 TapeOp::Div(l, r) => {
3638 let rv = vals[*r];
3640 adj[*l] += a / rv;
3641 adj[*r] -= a * vals[i] / rv;
3642 }
3643 TapeOp::Pow(l, r) => {
3644 let lv = vals[*l];
3645 let rv = vals[*r];
3646 if rv != 0.0 {
3647 adj[*l] += a * rv * lv.powf(rv - 1.0);
3648 }
3649 if lv > 0.0 {
3650 adj[*r] += a * vals[i] * lv.ln();
3651 }
3652 }
3653 TapeOp::Neg(j) => {
3654 adj[*j] -= a;
3655 }
3656 TapeOp::Abs(j) => {
3657 if vals[*j] >= 0.0 {
3658 adj[*j] += a;
3659 } else {
3660 adj[*j] -= a;
3661 }
3662 }
3663 TapeOp::Sqrt(j) => {
3664 let sv = vals[i];
3665 if sv > 0.0 {
3666 adj[*j] += a * 0.5 / sv;
3667 }
3668 }
3669 TapeOp::Exp(j) => {
3670 adj[*j] += a * vals[i];
3671 }
3672 TapeOp::Log(j) => {
3673 adj[*j] += a / vals[*j];
3674 }
3675 TapeOp::Log10(j) => {
3676 adj[*j] += div_u_ln10(a, vals[*j]);
3677 }
3678 TapeOp::Sin(j) => {
3679 adj[*j] += a * vals[*j].cos();
3680 }
3681 TapeOp::Cos(j) => {
3682 adj[*j] -= a * vals[*j].sin();
3683 }
3684 TapeOp::Tan(j) => {
3685 let t = vals[i];
3686 adj[*j] += a * (1.0 + t * t);
3687 }
3688 TapeOp::Atan(j) => {
3689 adj[*j] += a * atan_d1(vals[*j]);
3690 }
3691 TapeOp::Acos(j) => {
3692 adj[*j] += a * acos_d1(vals[*j]);
3693 }
3694 TapeOp::Sinh(j) => {
3695 adj[*j] += a * vals[*j].cosh();
3696 }
3697 TapeOp::Cosh(j) => {
3698 adj[*j] += a * vals[*j].sinh();
3699 }
3700 TapeOp::Tanh(j) => {
3701 adj[*j] += a * tanh_d1(vals[*j]);
3702 }
3703 TapeOp::Asin(j) => {
3704 adj[*j] += a * asin_d1(vals[*j]);
3705 }
3706 TapeOp::Acosh(j) => {
3707 adj[*j] += a * acosh_d1(vals[*j]);
3708 }
3709 TapeOp::Asinh(j) => {
3710 adj[*j] += a * asinh_d1(vals[*j]);
3711 }
3712 TapeOp::Atanh(j) => {
3713 adj[*j] += a * atanh_d1(vals[*j]);
3714 }
3715 TapeOp::Erf(j) => {
3716 adj[*j] += a * erf_d1(vals[*j]);
3717 }
3718 TapeOp::XLogX(j) => {
3719 adj[*j] += a * xlogx_d1(vals[*j]);
3720 }
3721 TapeOp::CEntropy(l, r) => {
3722 adj[*l] += a * centropy_da(vals[*l], vals[*r]);
3723 adj[*r] += a * centropy_db(vals[*l], vals[*r]);
3724 }
3725 TapeOp::Atan2(l, r) => {
3726 let y = vals[*l];
3727 let x = vals[*r];
3728 let (fa, fb) = atan2_d1(y, x);
3729 adj[*l] += a * fa;
3730 adj[*r] += a * fb;
3731 }
3732 TapeOp::Cmp(_, _, _)
3733 | TapeOp::And(_, _)
3734 | TapeOp::Or(_, _)
3735 | TapeOp::Not(_)
3736 | TapeOp::Select(_, _, _)
3737 | TapeOp::Min(_, _)
3738 | TapeOp::Max(_, _) => panic!(
3739 "GlobalTape free-function kernels do not implement conditional / logical \
3740 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3741 instead."
3742 ),
3743 TapeOp::Funcall(fc) => {
3744 let FuncallData { lib, name, args } = fc.as_ref();
3745 let call_args = funcall_to_ext_args(args, vals);
3746 let res = lib
3747 .eval(name, &call_args, true, false)
3748 .unwrap_or_else(|e| panic!("external function '{name}' reverse eval failed: {e}"));
3749 let derivs = res.derivs.expect("want_derivs=true returns derivs");
3750 let mut k = 0usize;
3751 for arg in args {
3752 if let TapeFuncallArg::Tape(idx) = arg {
3753 adj[*idx] += a * derivs[k];
3754 k += 1;
3755 }
3756 }
3757 let _ = i;
3758 let _ = grad;
3759 }
3760 }
3761}
3762
3763#[inline]
3771fn fwd_dir_step(op: &TapeOp, seed: &[f64], vals: &[f64], dot: &[f64], i: usize) -> f64 {
3772 match op {
3773 TapeOp::Const(_) => 0.0,
3774 TapeOp::Var(k) => seed[*k],
3775 TapeOp::Add(a, b) => dot[*a] + dot[*b],
3776 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
3777 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
3778 TapeOp::Div(a, b) => {
3779 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
3782 }
3783 TapeOp::Pow(a, b) => {
3784 let u = vals[*a];
3785 let r = vals[*b];
3786 let du = dot[*a];
3787 let dr = dot[*b];
3788 let mut result = 0.0;
3789 if r != 0.0 {
3794 result += r * u.powf(r - 1.0) * du;
3795 }
3796 if u > 0.0 {
3797 result += vals[i] * u.ln() * dr;
3798 }
3799 result
3800 }
3801 TapeOp::Neg(a) => -dot[*a],
3802 TapeOp::Abs(a) => {
3803 if vals[*a] >= 0.0 {
3804 dot[*a]
3805 } else {
3806 -dot[*a]
3807 }
3808 }
3809 TapeOp::Sqrt(a) => {
3810 let sv = vals[i];
3811 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
3812 }
3813 TapeOp::Exp(a) => vals[i] * dot[*a],
3814 TapeOp::Log(a) => dot[*a] / vals[*a],
3815 TapeOp::Log10(a) => div_u_ln10(dot[*a], vals[*a]),
3816 TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
3817 TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
3818 TapeOp::Tan(a) => {
3819 let t = vals[i];
3820 (1.0 + t * t) * dot[*a]
3821 }
3822 TapeOp::Atan(a) => dot[*a] * atan_d1(vals[*a]),
3823 TapeOp::Acos(a) => dot[*a] * acos_d1(vals[*a]),
3824 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
3825 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
3826 TapeOp::Tanh(a) => dot[*a] * tanh_d1(vals[*a]),
3827 TapeOp::Asin(a) => dot[*a] * asin_d1(vals[*a]),
3828 TapeOp::Acosh(a) => dot[*a] * acosh_d1(vals[*a]),
3829 TapeOp::Asinh(a) => dot[*a] * asinh_d1(vals[*a]),
3830 TapeOp::Atanh(a) => dot[*a] * atanh_d1(vals[*a]),
3831 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
3832 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
3833 TapeOp::CEntropy(a, b) => {
3834 centropy_da(vals[*a], vals[*b]) * dot[*a] + centropy_db(vals[*a], vals[*b]) * dot[*b]
3835 }
3836 TapeOp::Atan2(a, b) => {
3837 let y = vals[*a];
3838 let x = vals[*b];
3839 let (fa, fb) = atan2_d1(y, x);
3840 fa * dot[*a] + fb * dot[*b]
3841 }
3842 TapeOp::Cmp(_, _, _)
3843 | TapeOp::And(_, _)
3844 | TapeOp::Or(_, _)
3845 | TapeOp::Not(_)
3846 | TapeOp::Select(_, _, _)
3847 | TapeOp::Min(_, _)
3848 | TapeOp::Max(_, _) => panic!(
3849 "GlobalTape free-function kernels do not implement conditional / logical \
3850 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3851 instead."
3852 ),
3853 TapeOp::Funcall(fc) => {
3854 let FuncallData { lib, name, args } = fc.as_ref();
3855 let call_args = funcall_to_ext_args(args, vals);
3856 let res = lib
3857 .eval(name, &call_args, true, false)
3858 .unwrap_or_else(|e| panic!("external function '{name}' tangent eval failed: {e}"));
3859 let derivs = res.derivs.expect("want_derivs=true returns derivs");
3860 let mut acc = 0.0;
3861 let mut k = 0usize;
3862 for arg in args {
3863 if let TapeFuncallArg::Tape(idx) = arg {
3864 acc += derivs[k] * dot[*idx];
3865 k += 1;
3866 }
3867 }
3868 let _ = seed;
3869 acc
3870 }
3871 }
3872}
3873
3874#[allow(clippy::too_many_arguments)]
3883#[inline]
3884fn ror_dir_step(
3885 op: &TapeOp,
3886 i: usize,
3887 vals: &[f64],
3888 dot: &[f64],
3889 adj: &mut [f64],
3890 adj_dot: &mut [f64],
3891 w: f64,
3892 wd: f64,
3893 weight: f64,
3894 out: &mut [f64],
3895) {
3896 match op {
3897 TapeOp::Const(_) => {}
3898 TapeOp::Var(k) => {
3899 if wd != 0.0 {
3900 out[*k] += weight * wd;
3901 }
3902 }
3903 TapeOp::Add(a, b) => {
3904 adj[*a] += w;
3905 adj[*b] += w;
3906 adj_dot[*a] += wd;
3907 adj_dot[*b] += wd;
3908 }
3909 TapeOp::Sub(a, b) => {
3910 adj[*a] += w;
3911 adj[*b] -= w;
3912 adj_dot[*a] += wd;
3913 adj_dot[*b] -= wd;
3914 }
3915 TapeOp::Mul(a, b) => {
3916 adj[*a] += w * vals[*b];
3917 adj[*b] += w * vals[*a];
3918 adj_dot[*a] += wd * vals[*b] + w * dot[*b];
3919 adj_dot[*b] += wd * vals[*a] + w * dot[*a];
3920 }
3921 TapeOp::Div(a, b) => {
3922 let vb = vals[*b];
3925 let q = vals[i];
3926 let qd = dot[i];
3927 adj[*a] += w / vb;
3928 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
3929 adj[*b] -= w * q / vb;
3930 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
3931 }
3932 TapeOp::Pow(a, b) => {
3933 let u = vals[*a];
3934 let r = vals[*b];
3935 let du = dot[*a];
3936 let dr = dot[*b];
3937 if r != 0.0 {
3938 if u != 0.0 {
3939 let p_a = r * u.powf(r - 1.0);
3940 adj[*a] += w * p_a;
3941 let mut dp_a = dr * u.powf(r - 1.0);
3942 if u > 0.0 {
3943 dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
3944 } else {
3945 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
3946 }
3947 adj_dot[*a] += wd * p_a + w * dp_a;
3948 } else if r >= 2.0 {
3949 let p_a = 0.0;
3950 adj[*a] += w * p_a;
3951 let dp_a = if r == 2.0 {
3952 2.0 * du
3953 } else {
3954 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
3955 };
3956 adj_dot[*a] += wd * p_a + w * dp_a;
3957 }
3958 }
3959 if u > 0.0 {
3960 let ln_u = u.ln();
3961 let p_b = vals[i] * ln_u;
3962 adj[*b] += w * p_b;
3963 let dur = vals[i] * (r * du / u + dr * ln_u);
3964 let dp_b = dur * ln_u + vals[i] * du / u;
3965 adj_dot[*b] += wd * p_b + w * dp_b;
3966 }
3967 }
3968 TapeOp::Neg(a) => {
3969 adj[*a] -= w;
3970 adj_dot[*a] -= wd;
3971 }
3972 TapeOp::Abs(a) => {
3973 let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
3974 adj[*a] += w * s;
3975 adj_dot[*a] += wd * s;
3976 }
3977 TapeOp::Sqrt(a) => {
3978 let sv = vals[i];
3979 if sv > 0.0 {
3980 let fp = 0.5 / sv;
3981 let fpp = sqrt_d2(vals[*a], sv);
3982 adj[*a] += w * fp;
3983 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
3984 }
3985 }
3986 TapeOp::Exp(a) => {
3987 let ev = vals[i];
3988 adj[*a] += w * ev;
3989 adj_dot[*a] += wd * ev + w * ev * dot[*a];
3990 }
3991 TapeOp::Log(a) => {
3992 let u = vals[*a];
3993 adj[*a] += w / u;
3994 adj_dot[*a] += wd / u + w * log_d2(u) * dot[*a];
3995 }
3996 TapeOp::Log10(a) => {
3997 let u = vals[*a];
3998 adj[*a] += div_u_ln10(w, u);
3999 adj_dot[*a] += div_u_ln10(wd, u) + w * log10_d2(u) * dot[*a];
4000 }
4001 TapeOp::Sin(a) => {
4002 let u = vals[*a];
4003 let cu = u.cos();
4004 adj[*a] += w * cu;
4005 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
4006 }
4007 TapeOp::Cos(a) => {
4008 let u = vals[*a];
4009 let su = u.sin();
4010 adj[*a] -= w * su;
4011 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
4012 }
4013 TapeOp::Tan(a) => {
4014 let t = vals[i];
4015 let gp = 1.0 + t * t;
4016 let gpp = 2.0 * t * gp;
4017 adj[*a] += w * gp;
4018 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4019 }
4020 TapeOp::Atan(a) => {
4021 let u = vals[*a];
4022 let gp = atan_d1(u);
4023 let gpp = atan_d2(u);
4024 adj[*a] += w * gp;
4025 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4026 }
4027 TapeOp::Acos(a) => {
4028 let u = vals[*a];
4029 let gp = acos_d1(u);
4030 let gpp = acos_d2(u);
4031 adj[*a] += w * gp;
4032 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4033 }
4034 TapeOp::Sinh(a) => {
4035 let u = vals[*a];
4036 let gp = u.cosh();
4037 let gpp = vals[i]; adj[*a] += w * gp;
4039 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4040 }
4041 TapeOp::Cosh(a) => {
4042 let u = vals[*a];
4043 let gp = u.sinh();
4044 let gpp = vals[i]; adj[*a] += w * gp;
4046 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4047 }
4048 TapeOp::Tanh(a) => {
4049 let t = vals[i];
4050 let gp = tanh_d1(vals[*a]);
4051 let gpp = -2.0 * t * gp;
4052 adj[*a] += w * gp;
4053 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4054 }
4055 TapeOp::Asin(a) => {
4056 let u = vals[*a];
4057 let gp = asin_d1(u);
4058 let gpp = asin_d2(u);
4059 adj[*a] += w * gp;
4060 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4061 }
4062 TapeOp::Acosh(a) => {
4063 let u = vals[*a];
4064 let gp = acosh_d1(u);
4065 let gpp = acosh_d2(u);
4066 adj[*a] += w * gp;
4067 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4068 }
4069 TapeOp::Asinh(a) => {
4070 let u = vals[*a];
4071 let gp = asinh_d1(u);
4072 let gpp = asinh_d2(u);
4073 adj[*a] += w * gp;
4074 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4075 }
4076 TapeOp::Atanh(a) => {
4077 let u = vals[*a];
4078 let gp = atanh_d1(u);
4079 let gpp = atanh_d2(u);
4080 adj[*a] += w * gp;
4081 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4082 }
4083 TapeOp::Erf(a) => {
4084 let u = vals[*a];
4085 let gp = erf_d1(u);
4086 let gpp = erf_d2(u);
4087 adj[*a] += w * gp;
4088 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4089 }
4090 TapeOp::XLogX(a) => {
4091 let u = vals[*a];
4094 let gp = xlogx_d1(u);
4095 let gpp = xlogx_d2(u);
4096 adj[*a] += w * gp;
4097 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
4098 }
4099 TapeOp::CEntropy(a, b) => {
4100 let ua = vals[*a];
4101 let ub = vals[*b];
4102 let fa = centropy_da(ua, ub);
4103 let fb = centropy_db(ua, ub);
4104 let faa = centropy_daa(ua);
4105 let fab = centropy_dab(ub);
4106 let fbb = centropy_dbb(ua, ub);
4107 adj[*a] += w * fa;
4108 adj[*b] += w * fb;
4109 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
4110 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
4111 }
4112 TapeOp::Atan2(a, b) => {
4113 let y = vals[*a];
4114 let x = vals[*b];
4115 let (fa, fb) = atan2_d1(y, x);
4116 let (faa, fab, fbb) = atan2_d2(y, x);
4117 adj[*a] += w * fa;
4118 adj[*b] += w * fb;
4119 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
4120 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
4121 }
4122 TapeOp::Cmp(_, _, _)
4123 | TapeOp::And(_, _)
4124 | TapeOp::Or(_, _)
4125 | TapeOp::Not(_)
4126 | TapeOp::Select(_, _, _)
4127 | TapeOp::Min(_, _)
4128 | TapeOp::Max(_, _) => panic!(
4129 "GlobalTape free-function kernels do not implement conditional / logical \
4130 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
4131 instead."
4132 ),
4133 TapeOp::Funcall(fc) => {
4134 let FuncallData { lib, name, args } = fc.as_ref();
4135 let call_args = funcall_to_ext_args(args, vals);
4136 let res = lib.eval(name, &call_args, true, true).unwrap_or_else(|e| {
4137 panic!("external function '{name}' 2nd-order eval failed: {e}")
4138 });
4139 let derivs = res.derivs.expect("want_derivs=true returns derivs");
4140 let hes = res.hessian.expect("want_hes=true returns hessian");
4141 let real_tape: Vec<usize> = args
4142 .iter()
4143 .filter_map(|a| match a {
4144 TapeFuncallArg::Tape(t) => Some(*t),
4145 TapeFuncallArg::Str(_) => None,
4146 })
4147 .collect();
4148 for (k, &tk) in real_tape.iter().enumerate() {
4149 adj[tk] += w * derivs[k];
4150 let mut second_term = 0.0;
4151 for (l, &tl) in real_tape.iter().enumerate() {
4152 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
4153 let h_kl = hes[lo + hi * (hi + 1) / 2];
4154 second_term += h_kl * dot[tl];
4155 }
4156 adj_dot[tk] += wd * derivs[k] + w * second_term;
4157 }
4158 let _ = out;
4159 let _ = weight;
4160 let _ = i;
4161 }
4162 }
4163}
4164
4165fn hessian_sparsity_impl(ops: &[TapeOp]) -> BTreeSet<(usize, usize)> {
4169 let n = ops.len();
4170 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
4171 let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
4172
4173 let emit_cross =
4174 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4175 for &v1 in s1 {
4176 for &v2 in s2 {
4177 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
4178 pairs.insert((r, c));
4179 }
4180 }
4181 };
4182 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4183 let vars: Vec<usize> = s.iter().copied().collect();
4184 for (ai, &vi) in vars.iter().enumerate() {
4185 for &vj in &vars[..=ai] {
4186 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4187 pairs.insert((r, c));
4188 }
4189 }
4190 };
4191
4192 for op in ops {
4193 let vset = match op {
4194 TapeOp::Const(_) => BTreeSet::new(),
4195 TapeOp::Var(j) => {
4196 let mut s = BTreeSet::new();
4197 s.insert(*j);
4198 s
4199 }
4200 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
4201 var_sets[*a].union(&var_sets[*b]).copied().collect()
4202 }
4203 TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
4204 TapeOp::Mul(a, b) => {
4205 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4206 var_sets[*a].union(&var_sets[*b]).copied().collect()
4207 }
4208 TapeOp::Div(a, b) => {
4209 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4210 emit_self(&var_sets[*b], &mut pairs);
4211 var_sets[*a].union(&var_sets[*b]).copied().collect()
4212 }
4213 TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
4214 let combined: BTreeSet<usize> =
4215 var_sets[*a].union(&var_sets[*b]).copied().collect();
4216 emit_self(&combined, &mut pairs);
4217 combined
4218 }
4219 TapeOp::Sqrt(a)
4220 | TapeOp::Exp(a)
4221 | TapeOp::Log(a)
4222 | TapeOp::Log10(a)
4223 | TapeOp::Sin(a)
4224 | TapeOp::Cos(a)
4225 | TapeOp::Tan(a)
4226 | TapeOp::Atan(a)
4227 | TapeOp::Acos(a)
4228 | TapeOp::Sinh(a)
4229 | TapeOp::Cosh(a)
4230 | TapeOp::Tanh(a)
4231 | TapeOp::Asin(a)
4232 | TapeOp::Acosh(a)
4233 | TapeOp::Asinh(a)
4234 | TapeOp::Erf(a)
4235 | TapeOp::XLogX(a)
4236 | TapeOp::Atanh(a) => {
4237 emit_self(&var_sets[*a], &mut pairs);
4238 var_sets[*a].clone()
4239 }
4240 TapeOp::Funcall(fc) => {
4241 let args = &fc.args;
4242 let mut combined: BTreeSet<usize> = BTreeSet::new();
4243 for arg in args {
4244 if let TapeFuncallArg::Tape(t) = arg {
4245 for &vv in &var_sets[*t] {
4246 combined.insert(vv);
4247 }
4248 }
4249 }
4250 emit_self(&combined, &mut pairs);
4251 combined
4252 }
4253 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
4254 BTreeSet::new()
4257 }
4258 TapeOp::Select(_, t, e) => {
4259 var_sets[*t].union(&var_sets[*e]).copied().collect()
4262 }
4263 TapeOp::Min(a, b) | TapeOp::Max(a, b) => {
4264 var_sets[*a].union(&var_sets[*b]).copied().collect()
4267 }
4268 };
4269 var_sets.push(vset);
4270 }
4271 pairs
4272}
4273
4274#[cfg(test)]
4275mod tests {
4276 use super::*;
4277
4278 fn cnst(c: f64) -> Expr {
4279 Expr::Const(c)
4280 }
4281 fn var(i: usize) -> Expr {
4282 Expr::Var(i)
4283 }
4284 fn add(a: Expr, b: Expr) -> Expr {
4285 Expr::Binary(BinOp::Add, Box::new(a), Box::new(b))
4286 }
4287 fn mul(a: Expr, b: Expr) -> Expr {
4288 Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b))
4289 }
4290 fn pow(a: Expr, b: Expr) -> Expr {
4291 Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b))
4292 }
4293 fn div(a: Expr, b: Expr) -> Expr {
4294 Expr::Binary(BinOp::Div, Box::new(a), Box::new(b))
4295 }
4296 fn unary(op: UnaryOp, a: Expr) -> Expr {
4297 Expr::Unary(op, Box::new(a))
4298 }
4299 fn cmp(op: CmpOp, a: Expr, b: Expr) -> Expr {
4300 Expr::Compare(op, Box::new(a), Box::new(b))
4301 }
4302 fn cond(c: Expr, t: Expr, e: Expr) -> Expr {
4303 Expr::Cond {
4304 cond: Box::new(c),
4305 then_: Box::new(t),
4306 else_: Box::new(e),
4307 }
4308 }
4309
4310 #[test]
4311 fn polynomial_eval_and_grad() {
4312 let e = add(
4314 mul(cnst(3.0), pow(var(0), cnst(2.0))),
4315 mul(cnst(2.0), var(1)),
4316 );
4317 let t = Tape::build(&e);
4318 assert!((t.eval(&[2.0, 3.0]) - 18.0).abs() < 1e-12);
4319 let mut g = vec![0.0; 2];
4320 t.gradient_seed(&[2.0, 3.0], 1.0, &mut g);
4321 assert!((g[0] - 12.0).abs() < 1e-12);
4323 assert!((g[1] - 2.0).abs() < 1e-12);
4324 }
4325
4326 #[test]
4327 fn cse_shared_body_evaluated_once() {
4328 let body = Arc::new(add(var(0), var(1)));
4330 let e = add(
4331 pow(Expr::Cse(body.clone()), cnst(2.0)),
4332 Expr::Cse(body.clone()),
4333 );
4334 let t = Tape::build(&e);
4335 let n_body_adds = t
4337 .ops
4338 .iter()
4339 .filter(|op| {
4340 matches!(op, TapeOp::Add(a, b) if {
4341 matches!(t.ops[*a], TapeOp::Var(0)) && matches!(t.ops[*b], TapeOp::Var(1))
4342 })
4343 })
4344 .count();
4345 assert_eq!(n_body_adds, 1, "CSE body should be emitted exactly once");
4346
4347 assert!((t.eval(&[1.0, 2.0]) - 12.0).abs() < 1e-12);
4349 let mut g = vec![0.0; 2];
4350 t.gradient_seed(&[1.0, 2.0], 1.0, &mut g);
4351 assert!((g[0] - 7.0).abs() < 1e-12);
4353 assert!((g[1] - 7.0).abs() < 1e-12);
4354 }
4355
4356 fn fd_check(tape: &Tape, x: &[f64], n: usize, tol: f64) {
4357 let vars = tape.variables();
4358 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4359 let mut pairs = Vec::new();
4360 for (ai, &vi) in vars.iter().enumerate() {
4361 for &vj in &vars[..=ai] {
4362 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4363 hess_map.entry((r, c)).or_insert_with(|| {
4364 let p = pairs.len();
4365 pairs.push((r, c));
4366 p
4367 });
4368 }
4369 }
4370 let nnz = pairs.len();
4371 let mut ad = vec![0.0; nnz];
4372 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
4373
4374 let mut fd = vec![0.0; nnz];
4375 let mut xp = x.to_vec();
4376 let mut gp = vec![0.0; n];
4377 let mut gm = vec![0.0; n];
4378 for &j in &vars {
4379 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4380 xp[j] = x[j] + h;
4381 gp.iter_mut().for_each(|v| *v = 0.0);
4382 tape.gradient_seed(&xp, 1.0, &mut gp);
4383 xp[j] = x[j] - h;
4384 gm.iter_mut().for_each(|v| *v = 0.0);
4385 tape.gradient_seed(&xp, 1.0, &mut gm);
4386 xp[j] = x[j];
4387 for &i in &vars {
4388 if i >= j {
4389 if let Some(&pos) = hess_map.get(&(i, j)) {
4390 fd[pos] = (gp[i] - gm[i]) / (2.0 * h);
4391 }
4392 }
4393 }
4394 }
4395 for (k, &(r, c)) in pairs.iter().enumerate() {
4396 let scale = fd[k].abs().max(1.0);
4397 assert!(
4398 (ad[k] - fd[k]).abs() / scale < tol,
4399 "H[{},{}]: AD={:.6e} FD={:.6e}",
4400 r,
4401 c,
4402 ad[k],
4403 fd[k]
4404 );
4405 }
4406 }
4407
4408 #[test]
4409 fn hessian_quadratic_matches_fd() {
4410 let e = add(
4412 add(
4413 mul(cnst(3.0), pow(var(0), cnst(2.0))),
4414 mul(cnst(2.0), mul(var(0), var(1))),
4415 ),
4416 pow(var(1), cnst(2.0)),
4417 );
4418 let t = Tape::build(&e);
4419 fd_check(&t, &[2.0, 3.0], 2, 1e-5);
4420 }
4421
4422 #[test]
4423 fn hessian_transcendental_matches_fd() {
4424 let e = Expr::Sum(vec![
4426 unary(UnaryOp::Exp, var(0)),
4427 unary(UnaryOp::Sin, var(1)),
4428 unary(UnaryOp::Log, var(0)),
4429 unary(UnaryOp::Sqrt, var(1)),
4430 mul(var(0), var(1)),
4431 ]);
4432 let t = Tape::build(&e);
4433 fd_check(&t, &[1.5, 2.0], 2, 1e-5);
4434 }
4435
4436 #[test]
4437 fn inverse_trig_grad_and_hessian_match_fd() {
4438 let e = Expr::Sum(vec![
4442 unary(UnaryOp::Tan, var(0)),
4443 unary(UnaryOp::Atan, var(1)),
4444 unary(UnaryOp::Acos, var(2)),
4445 mul(var(0), var(1)),
4446 ]);
4447 let t = Tape::build(&e);
4448 let x = [0.5, 1.3, 0.3];
4449
4450 let mut g = vec![0.0; 3];
4454 t.gradient_seed(&x, 1.0, &mut g);
4455 for j in 0..3 {
4456 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4457 let mut xp = x;
4458 let mut xm = x;
4459 xp[j] += h;
4460 xm[j] -= h;
4461 let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4462 let scale = fd.abs().max(1.0);
4463 assert!(
4464 (g[j] - fd).abs() / scale < 1e-5,
4465 "grad[{j}]: AD={:.6e} FD={:.6e}",
4466 g[j],
4467 fd
4468 );
4469 }
4470
4471 fd_check(&t, &x, 3, 1e-5);
4473 }
4474
4475 fn grad_and_hess_match_fd(e: &Expr, x: &[f64], tol: f64) {
4478 let n = x.len();
4479 let t = Tape::build(e);
4480 let mut g = vec![0.0; n];
4481 t.gradient_seed(x, 1.0, &mut g);
4482 for j in 0..n {
4483 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4484 let mut xp = x.to_vec();
4485 let mut xm = x.to_vec();
4486 xp[j] += h;
4487 xm[j] -= h;
4488 let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4489 let scale = fd.abs().max(1.0);
4490 assert!(
4491 (g[j] - fd).abs() / scale < tol,
4492 "grad[{j}]: AD={:.6e} FD={:.6e}",
4493 g[j],
4494 fd
4495 );
4496 }
4497 fd_check(&t, x, n, tol);
4498 }
4499
4500 #[test]
4501 fn hyperbolic_grad_and_hessian_match_fd() {
4502 let e = Expr::Sum(vec![
4505 unary(UnaryOp::Sinh, var(0)),
4506 unary(UnaryOp::Cosh, var(1)),
4507 unary(UnaryOp::Tanh, var(2)),
4508 unary(UnaryOp::Asinh, var(3)),
4509 mul(var(0), var(1)),
4510 mul(var(2), var(3)),
4511 ]);
4512 grad_and_hess_match_fd(&e, &[0.5, 0.7, 0.3, 1.1], 1e-5);
4513 }
4514
4515 #[test]
4516 fn restricted_inverse_grad_and_hessian_match_fd() {
4517 let e = Expr::Sum(vec![
4521 unary(UnaryOp::Asin, var(0)),
4522 unary(UnaryOp::Acosh, var(1)),
4523 unary(UnaryOp::Atanh, var(2)),
4524 mul(var(0), var(2)),
4525 ]);
4526 grad_and_hess_match_fd(&e, &[0.4, 1.8, 0.3], 1e-5);
4527 }
4528
4529 #[test]
4530 fn erf_value_matches_reference() {
4531 let t = Tape::build(&unary(UnaryOp::Erf, var(0)));
4536 for (x, want) in [
4537 (0.0, 0.0),
4538 (0.5, 0.520_499_877_813_046_5),
4539 (1.0, 0.842_700_792_949_714_9),
4540 (-1.0, -0.842_700_792_949_714_9),
4541 (2.0, 0.995_322_265_018_952_7),
4542 (3.0, 0.999_977_909_503_001_4),
4543 ] {
4544 let got = t.eval(&[x]);
4545 assert!(
4546 (got - want).abs() < 1e-15,
4547 "erf({x}): got {got:.17e}, want {want:.17e}"
4548 );
4549 }
4550 assert!((t.eval(&[-0.3]) + t.eval(&[0.3])).abs() < 1e-16);
4553 assert!((t.eval(&[10.0]) - 1.0).abs() < 1e-15);
4554 }
4555
4556 #[test]
4557 fn atanh_value_and_derivative_accurate_near_boundary() {
4558 let t = Tape::build(&unary(UnaryOp::Atanh, var(0)));
4568 let rel = |got: f64, want: f64| ((got - want) / want).abs();
4569 for (u, want_v, want_d) in [
4570 (-0.9999999999999, -15.31322124631433, 4998445757348.192),
4571 (-0.999999999999999, -17.616361586450413, 500399958596722.0),
4572 (0.999999999999999, 17.616361586450413, 500399958596722.0),
4573 (0.9999999999999998, 18.36840028483855, 2251799813685248.5),
4574 (-0.9, -1.4722194895832204, 5.263157894736843),
4575 (0.5, 0.5493061443340549, 1.3333333333333333),
4576 (1e-10, 1e-10, 1.0),
4577 ] {
4578 let got_v = t.eval(&[u]);
4579 assert!(
4580 rel(got_v, want_v) < 4e-16,
4581 "atanh({u:e}): got {got_v:.17e}, want {want_v:.17e}"
4582 );
4583 let mut g = [0.0];
4584 t.gradient_seed(&[u], 1.0, &mut g);
4585 assert!(
4586 rel(g[0], want_d) < 4e-16,
4587 "atanh'({u:e}): got {:.17e}, want {want_d:.17e}",
4588 g[0]
4589 );
4590 }
4591 }
4592
4593 fn close(got: f64, want: f64, tol: f64) -> bool {
4597 if want == 0.0 {
4598 got == 0.0
4599 } else if want.abs() < f64::MIN_POSITIVE {
4600 (got - want).abs() <= f64::from_bits(1)
4601 } else {
4602 ((got - want) / want).abs() < tol
4603 }
4604 }
4605
4606 const INVERSE_DERIV_CASES: &[(UnaryOp, f64, f64, f64)] = &[
4617 (UnaryOp::Asin, 0.3, 1.0482848367219182, 0.3455884077105225),
4618 (
4619 UnaryOp::Asin,
4620 -0.999999999999999,
4621 22369621.33333334,
4622 -1.1193757589024341e22,
4623 ),
4624 (
4625 UnaryOp::Asin,
4626 0.999999999999999,
4627 22369621.33333334,
4628 1.1193757589024341e22,
4629 ),
4630 (
4631 UnaryOp::Asin,
4632 0.999999992554274,
4633 8194.670143394766,
4634 550293555898.1869,
4635 ),
4636 (UnaryOp::Acos, 0.3, -1.0482848367219182, -0.3455884077105225),
4637 (
4638 UnaryOp::Acos,
4639 -0.999999999999999,
4640 -22369621.33333334,
4641 1.1193757589024341e22,
4642 ),
4643 (
4644 UnaryOp::Acos,
4645 0.999999992554274,
4646 -8194.670143394766,
4647 -550293555898.1869,
4648 ),
4649 (UnaryOp::Acosh, 2.0, 0.5773502691896257, -0.3849001794597505),
4650 (
4651 UnaryOp::Acosh,
4652 1.0000000000001,
4653 2236962.1333332774,
4654 -1.1193757589024625e19,
4655 ),
4656 (
4657 UnaryOp::Acosh,
4658 1.0000000050104683,
4659 9989.54813402654,
4660 -996867721306.0283,
4661 ),
4662 (UnaryOp::Acosh, 1e100, 1e-100, -1e-200),
4663 (UnaryOp::Acosh, 1e200, 1e-200, 0.0),
4664 (
4665 UnaryOp::Asinh,
4666 0.5,
4667 0.8944271909999159,
4668 -0.35777087639996635,
4669 ),
4670 (UnaryOp::Asinh, 1e100, 1e-100, -1e-200),
4671 (UnaryOp::Asinh, -1e100, 1e-100, 1e-200),
4672 (UnaryOp::Asinh, 1e200, 1e-200, 0.0),
4673 (UnaryOp::Asinh, -1e200, 1e-200, 0.0),
4674 (UnaryOp::Atanh, 0.5, 1.3333333333333333, 1.7777777777777777),
4675 (
4676 UnaryOp::Atanh,
4677 -0.999999999999999,
4678 500399958596722.0,
4679 -5.008002371272017e29,
4680 ),
4681 (
4682 UnaryOp::Atanh,
4683 0.999999999999999,
4684 500399958596722.0,
4685 5.008002371272017e29,
4686 ),
4687 (
4688 UnaryOp::Atanh,
4689 -0.999999992554274,
4690 67152618.75904559,
4691 -9018948345242822.0,
4692 ),
4693 (UnaryOp::Atan, 0.7, 0.6711409395973155, -0.6306022251249944),
4694 (UnaryOp::Atan, 1e100, 1e-200, -2e-300),
4695 (UnaryOp::Atan, -1e100, 1e-200, 2e-300),
4696 (UnaryOp::Atan, 2e150, 2.5e-301, 0.0),
4700 (UnaryOp::Atan, 1e155, 1e-310, 0.0),
4701 (UnaryOp::Atan, -1e158, 1e-316, 0.0),
4702 (UnaryOp::Atan, 1e160, 1e-320, 0.0),
4703 (UnaryOp::Atan, 5e161, 5e-324, 0.0),
4704 (UnaryOp::Tanh, 0.3, 0.9151369618266292, -0.5331818782014544),
4705 (
4706 UnaryOp::Tanh,
4707 10.0,
4708 8.244614455767397e-9,
4709 -1.6489228843561127e-8,
4710 ),
4711 (
4712 UnaryOp::Tanh,
4713 25.0,
4714 7.714999391855671e-22,
4715 -1.5429998783711343e-21,
4716 ),
4717 (
4718 UnaryOp::Tanh,
4719 -30.0,
4720 3.502604305078608e-26,
4721 7.005208610157216e-26,
4722 ),
4723 ];
4724
4725 #[test]
4726 fn inverse_and_tanh_derivatives_accurate_at_boundaries_and_extremes() {
4727 assert_unary_derivs(INVERSE_DERIV_CASES);
4735 }
4736
4737 fn assert_unary_derivs(cases: &[(UnaryOp, f64, f64, f64)]) {
4740 for &(op, u, want_d1, want_d2) in cases {
4741 let t = Tape::build(&unary(op, var(0)));
4742 let mut g = [0.0];
4743 t.gradient_seed(&[u], 1.0, &mut g);
4744 assert!(
4745 close(g[0], want_d1, 1e-15),
4746 "{op:?}'({u:e}): got {:.17e}, want {want_d1:.17e}",
4747 g[0]
4748 );
4749 let mut hess_map = HashMap::new();
4754 hess_map.insert((0usize, 0usize), 0usize);
4755 let mut h = [0.0];
4756 t.hessian_accumulate(&[u], 1.0, &hess_map, &mut h);
4757 assert!(
4758 close(h[0], want_d2, 2e-15),
4759 "{op:?}''({u:e}): got {:.17e}, want {want_d2:.17e}",
4760 h[0]
4761 );
4762 }
4763 }
4764
4765 const LOG_SQRT_DERIV_CASES: &[(UnaryOp, f64, f64, f64)] = &[
4772 (UnaryOp::Log, 2.0, 0.5, -0.25),
4773 (UnaryOp::Log, 1e100, 1e-100, -1e-200),
4774 (UnaryOp::Log, 1e155, 1e-155, -1e-310),
4775 (UnaryOp::Log, 1e160, 1e-160, -1e-320),
4776 (
4777 UnaryOp::Log10,
4778 2.0,
4779 0.2171472409516259,
4780 -0.10857362047581295,
4781 ),
4782 (
4783 UnaryOp::Log10,
4784 1e100,
4785 4.342944819032518e-101,
4786 -4.342944819032518e-201,
4787 ),
4788 (
4789 UnaryOp::Log10,
4790 1e155,
4791 4.3429448190325184e-156,
4792 -4.3429448190326e-311,
4793 ),
4794 (UnaryOp::Log10, 1e305, 4.3429448190325186e-306, 0.0),
4795 (UnaryOp::Log10, 1.7e308, 2.554673422960305e-309, 0.0),
4796 (UnaryOp::Sqrt, 2.0, 0.3535533905932738, -0.08838834764831845),
4797 (UnaryOp::Sqrt, 1e150, 5e-76, -2.5e-226),
4798 (UnaryOp::Sqrt, 1e206, 5e-104, -2.5e-310),
4799 (UnaryOp::Sqrt, 1e300, 5e-151, 0.0),
4800 (UnaryOp::Sqrt, 1.7e308, 3.8348249442368524e-155, 0.0),
4801 ];
4802
4803 #[test]
4804 fn log_log10_sqrt_derivatives_accurate_at_large_arguments() {
4805 assert_unary_derivs(LOG_SQRT_DERIV_CASES);
4806 }
4807
4808 #[test]
4809 fn atan2_derivatives_accurate_at_extreme_magnitudes() {
4810 let e = Expr::Binary(BinOp::Atan2, Box::new(var(0)), Box::new(var(1)));
4815 let t = Tape::build(&e);
4816 let mut hess_map = HashMap::new();
4817 hess_map.insert((0usize, 0usize), 0usize);
4818 hess_map.insert((1usize, 0usize), 1usize);
4819 hess_map.insert((1usize, 1usize), 2usize);
4820 for (y, x, want_g, want_h) in [
4821 (
4822 1e100,
4823 2e100,
4824 [4e-101, -2e-101],
4825 [-1.6e-201, -1.2e-201, 1.6e-201],
4826 ),
4827 (1e-100, -3e-100, [-3e99, -1e99], [6e198, -8e198, -6e198]),
4828 (1.0, 2.0, [0.4, -0.2], [-0.16, -0.12, 0.16]),
4829 (1e308, -1e308, [-5e-309, -5e-309], [0.0, 0.0, 0.0]),
4833 (
4834 1.5e308,
4835 -1e308,
4836 [-3.076923076923076e-309, -4.615384615384614e-309],
4837 [0.0, 0.0, 0.0],
4838 ),
4839 (1e308, 1e308, [5e-309, -5e-309], [0.0, 0.0, 0.0]),
4840 (
4841 3.0,
4842 3.0000000001,
4843 [0.16666666666666666, -0.1666666666611111],
4844 [
4845 -0.0555555555537037,
4846 -1.8518520049821686e-12,
4847 0.0555555555537037,
4848 ],
4849 ),
4850 ] {
4851 let mut g = [0.0; 2];
4852 t.gradient_seed(&[y, x], 1.0, &mut g);
4853 let mut h = [0.0; 3];
4854 t.hessian_accumulate(&[y, x], 1.0, &hess_map, &mut h);
4855 for k in 0..2 {
4856 assert!(
4857 close(g[k], want_g[k], 1e-15),
4858 "atan2({y:e},{x:e}) grad[{k}]: got {:.17e}, want {:.17e}",
4859 g[k],
4860 want_g[k]
4861 );
4862 }
4863 for k in 0..3 {
4864 assert!(
4865 close(h[k], want_h[k], 2e-15),
4866 "atan2({y:e},{x:e}) hess[{k}]: got {:.17e}, want {:.17e}",
4867 h[k],
4868 want_h[k]
4869 );
4870 }
4871 }
4872 }
4873
4874 #[test]
4875 fn erf_second_derivative_stays_finite_at_extreme_magnitudes() {
4876 for u in [1e19, 1e150, 1e300, f64::MAX, f64::MAX / 2.0] {
4884 for signed in [u, -u] {
4885 let d2 = erf_d2(signed);
4886 assert!(
4887 d2.is_finite(),
4888 "erf_d2({signed:e}) = {d2} — must be finite (the limit is 0)"
4889 );
4890 }
4891 }
4892 let u = 0.7;
4894 let want = -2.0 * u * (2.0 / std::f64::consts::PI.sqrt()) * (-u * u).exp();
4895 assert!((erf_d2(u) - want).abs() < 1e-15, "{}", erf_d2(u));
4896 }
4897
4898 #[test]
4899 fn quotient_rule_keeps_representable_derivatives_at_extreme_denominators() {
4900 let t = Tape::build(&div(var(0), var(1)));
4911 let mut g = vec![0.0; 2];
4912 for b in [1e300_f64, 1e200, 1e160, -1e300, 1e-160, 1e-300] {
4913 g.iter_mut().for_each(|v| *v = 0.0);
4919 t.gradient_seed(&[b, b], 1.0, &mut g);
4920 assert_eq!(g[0], 1.0 / b, "d(a/b)/da at a=b={b:e}");
4921 assert_eq!(
4922 g[1],
4923 -1.0 / b,
4924 "d(a/b)/db at a=b={b:e} — b² is not representable"
4925 );
4926 }
4927
4928 let b = 1e155_f64;
4932 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4933 hess_map.insert((0, 0), 0);
4934 hess_map.insert((1, 0), 1);
4935 hess_map.insert((1, 1), 2);
4936 let mut ad = vec![0.0; 3];
4937 t.hessian_accumulate(&[b, b], 1.0, &hess_map, &mut ad);
4938 let want = -(1.0 / b) / b;
4942 assert!(
4943 want != 0.0 && want.is_finite(),
4944 "test premise: -1/b² must be representable, got {want:e}"
4945 );
4946 assert!(
4947 (ad[1] - want).abs() <= 1e-10 * want.abs(),
4948 "d²(a/b)/da db at a=b={b:e}: got {:e}, want {want:e}",
4949 ad[1]
4950 );
4951 }
4952
4953 #[test]
4954 fn erf_grad_and_hessian_match_fd() {
4955 let e = Expr::Sum(vec![
4959 unary(UnaryOp::Erf, var(0)),
4960 unary(UnaryOp::Erf, mul(cnst(2.0), var(1))),
4961 mul(var(0), var(1)),
4962 ]);
4963 grad_and_hess_match_fd(&e, &[0.4, -0.7], 1e-5);
4964 }
4965
4966 #[test]
4967 fn erf_directional_hessian_matches_accumulated() {
4968 let e = Expr::Sum(vec![
4973 unary(UnaryOp::Erf, var(0)),
4974 unary(UnaryOp::Erf, mul(var(0), var(1))),
4975 ]);
4976 let tape = Tape::build(&e);
4977 let x = [0.6, -0.9];
4978 let n = x.len();
4979
4980 let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
4981 let hess_map: HashMap<(usize, usize), usize> =
4982 pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
4983 let mut acc = vec![0.0; pairs.len()];
4984 tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
4985
4986 let ops = tape.ops.len();
4988 let mut vals = vec![0.0; ops];
4989 tape.forward_into(&x, &mut vals);
4990 for j in 0..n {
4991 let mut seed = vec![0.0; n];
4992 seed[j] = 1.0;
4993 let mut col = vec![0.0; n];
4994 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
4995 tape.hessian_directional(
4996 &vals,
4997 &seed,
4998 1.0,
4999 &mut col,
5000 &mut dot,
5001 &mut adj,
5002 &mut adj_dot,
5003 );
5004 for i in 0..n {
5005 let (r, c) = if i >= j { (i, j) } else { (j, i) };
5006 let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
5007 assert!(
5008 (col[i] - want).abs() < 1e-12,
5009 "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
5010 col[i]
5011 );
5012 }
5013 }
5014 }
5015
5016 fn centropy_expr(a: Expr, b: Expr) -> Expr {
5017 Expr::Binary(BinOp::CEntropy, Box::new(a), Box::new(b))
5018 }
5019
5020 fn hess2(e: &Expr, x: &[f64; 2]) -> (f64, f64, f64) {
5022 let tape = Tape::build(e);
5023 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5024 hess_map.insert((0, 0), 0);
5025 hess_map.insert((1, 0), 1);
5026 hess_map.insert((1, 1), 2);
5027 let mut ad = vec![0.0; 3];
5028 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
5029 (ad[0], ad[1], ad[2])
5030 }
5031
5032 #[test]
5033 fn xlogx_and_centropy_grad_and_hessian_match_fd() {
5034 let e = Expr::Sum(vec![
5038 unary(UnaryOp::XLogX, var(0)),
5039 centropy_expr(var(0), var(1)),
5040 mul(var(0), var(1)),
5041 ]);
5042 grad_and_hess_match_fd(&e, &[1.7, 0.6], 1e-5);
5043 }
5044
5045 #[test]
5046 fn centropy_directional_hessian_matches_accumulated() {
5047 let e = Expr::Sum(vec![
5051 centropy_expr(var(0), var(1)),
5052 unary(UnaryOp::XLogX, mul(var(0), var(1))),
5053 ]);
5054 let tape = Tape::build(&e);
5055 let x = [1.3, 0.8];
5056 let n = x.len();
5057
5058 let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
5059 let hess_map: HashMap<(usize, usize), usize> =
5060 pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
5061 let mut acc = vec![0.0; pairs.len()];
5062 tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
5063
5064 let ops = tape.ops.len();
5065 let mut vals = vec![0.0; ops];
5066 tape.forward_into(&x, &mut vals);
5067 let mut compared = 0usize;
5068 for j in 0..n {
5069 let mut seed = vec![0.0; n];
5070 seed[j] = 1.0;
5071 let mut col = vec![0.0; n];
5072 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
5073 tape.hessian_directional(
5074 &vals,
5075 &seed,
5076 1.0,
5077 &mut col,
5078 &mut dot,
5079 &mut adj,
5080 &mut adj_dot,
5081 );
5082 for i in 0..n {
5083 let (r, c) = if i >= j { (i, j) } else { (j, i) };
5084 let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
5085 assert!(
5086 (col[i] - want).abs() < 1e-12,
5087 "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
5088 col[i]
5089 );
5090 compared += 1;
5091 }
5092 }
5093 assert_eq!(compared, 4, "expected a full 2x2 comparison");
5096 assert!(
5097 acc.iter().any(|v| v.abs() > 1e-12),
5098 "every accumulated Hessian entry was zero — the arms under test never ran"
5099 );
5100 }
5101
5102 #[test]
5103 fn fused_entropy_ops_reach_derivatives_the_chain_rule_cannot() {
5104 let a = 1e-299_f64;
5110 let want = 1.0 / a; assert!(want.is_finite(), "test premise: 1/a must be representable");
5112
5113 let (fused, _, _) = hess2(&unary(UnaryOp::XLogX, var(0)), &[a, 0.0]);
5114 assert!(
5115 (fused - want).abs() <= 1e-12 * want,
5116 "xlogx''({a:e}): got {fused:e}, want {want:e}"
5117 );
5118
5119 let (decomposed, _, _) = hess2(&mul(var(0), unary(UnaryOp::Log, var(0))), &[a, 0.0]);
5124 assert!(
5125 !decomposed.is_finite(),
5126 "test premise: a*log(a) is supposed to lose this second derivative, \
5127 but it returned {decomposed:e} — if the Log/Mul arms now reach it, \
5128 re-derive whether the fused op is still needed"
5129 );
5130
5131 let (fused_aa, _, _) = hess2(¢ropy_expr(var(0), var(1)), &[a, 1.0]);
5133 assert!(
5134 (fused_aa - want).abs() <= 1e-12 * want,
5135 "centropy ∂²/∂a² at a={a:e}: got {fused_aa:e}, want {want:e}"
5136 );
5137 }
5138
5139 #[test]
5140 fn centropy_second_derivative_in_b_survives_a_squared_denominator() {
5141 let (a, b) = (1e300_f64, 1e200_f64);
5144 let want = (a / b) / b;
5145 assert!(
5149 want.is_finite() && want != 0.0 && a / (b * b) == 0.0,
5150 "test premise: a/b² representable ({want:e}) but a/(b*b) collapses"
5151 );
5152 let (_, _, h11) = hess2(¢ropy_expr(var(0), var(1)), &[a, b]);
5153 assert!(
5154 (h11 - want).abs() <= 1e-12 * want,
5155 "centropy ∂²/∂b² at ({a:e}, {b:e}): got {h11:e}, want {want:e}"
5156 );
5157 }
5158
5159 #[test]
5160 fn centropy_value_survives_an_out_of_range_ratio_and_the_zero_limit() {
5161 let t = |e: &Expr, x: &[f64]| Tape::build(e).eval(x);
5162 let ce = centropy_expr(var(0), var(1));
5163
5164 let want = 1e300 * (1e300_f64.ln() - 1e-300_f64.ln());
5168 let got = t(&ce, &[1e300, 1e-300]);
5169 assert!(
5170 got.is_finite() && (got - want).abs() <= 1e-12 * want,
5171 "centropy(1e300, 1e-300): got {got:e}, want {want:e}"
5172 );
5173 assert!(
5175 !t(
5176 &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
5177 &[1e300, 1e-300]
5178 )
5179 .is_finite(),
5180 "test premise: a*log(a/b) is supposed to overflow here"
5181 );
5182
5183 let (a, b) = (1e10 + 1.0, 1e10_f64);
5190 assert_eq!(a - b, 1.0, "test premise: a - b must be exact here");
5191 let got = t(&ce, &[a, b]);
5192 assert!(
5193 (got - 1.00000000005).abs() <= 2.3e-16,
5194 "centropy(1e10+1, 1e10): got {got:.17e}, want 1.00000000005"
5195 );
5196 let naive = t(
5198 &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
5199 &[a, b],
5200 );
5201 assert!(
5202 (naive - 1.00000000005).abs() > 1e-9,
5203 "test premise: a*log(a/b) is supposed to lose digits here, got {naive:.17e}"
5204 );
5205
5206 assert_eq!(t(&ce, &[0.0, 2.0]), 0.0);
5208 assert_eq!(t(&unary(UnaryOp::XLogX, var(0)), &[0.0]), 0.0);
5209 }
5210
5211 #[test]
5212 fn cond_does_not_leak_a_non_finite_from_its_inactive_branch() {
5213 let cond = Expr::Compare(CmpOp::Lt, Box::new(var(0)), Box::new(cnst(1e-300)));
5221 let e = Expr::Cond {
5222 cond: Box::new(cond),
5223 then_: Box::new(mul(var(0), cnst(1e-300f64.ln()))),
5224 else_: Box::new(unary(UnaryOp::XLogX, var(0))),
5225 };
5226 let t = Tape::build(&e);
5227 let x = [-1.0_f64];
5228
5229 assert!(
5232 Tape::build(&unary(UnaryOp::XLogX, var(0)))
5233 .eval(&x)
5234 .is_nan(),
5235 "test premise: xlogx(-1) must be NaN"
5236 );
5237
5238 assert_eq!(t.eval(&x), -1e-300_f64.ln());
5239
5240 let mut g = vec![0.0; 1];
5241 t.gradient_seed(&x, 1.0, &mut g);
5242 assert_eq!(g[0], 1e-300_f64.ln(), "gradient leaked the inactive branch");
5243
5244 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5245 hess_map.insert((0, 0), 0);
5246 let mut acc = vec![0.0; 1];
5247 t.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
5248 assert_eq!(acc[0], 0.0, "hessian_accumulate leaked the inactive branch");
5249
5250 let ops = t.ops.len();
5251 let mut vals = vec![0.0; ops];
5252 t.forward_into(&x, &mut vals);
5253 let mut col = vec![0.0; 1];
5254 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
5255 t.hessian_directional(
5256 &vals,
5257 &[1.0],
5258 1.0,
5259 &mut col,
5260 &mut dot,
5261 &mut adj,
5262 &mut adj_dot,
5263 );
5264 assert_eq!(
5265 col[0], 0.0,
5266 "hessian_directional leaked the inactive branch"
5267 );
5268 }
5269
5270 #[test]
5271 fn atan2_grad_and_hessian_match_fd() {
5272 let atan2 = |a: Expr, b: Expr| Expr::Binary(BinOp::Atan2, Box::new(a), Box::new(b));
5274 let e = Expr::Sum(vec![atan2(var(0), var(1)), mul(var(0), var(1))]);
5275 grad_and_hess_match_fd(&e, &[1.2, 0.7], 1e-5);
5276 }
5277
5278 #[test]
5279 fn minmax_grad_and_hessian_match_fd() {
5280 let e = Expr::Sum(vec![
5287 Expr::MinList(vec![var(0), var(1), var(2)]),
5288 Expr::MaxList(vec![var(1), var(2)]),
5289 mul(var(0), var(2)),
5290 ]);
5291 grad_and_hess_match_fd(&e, &[0.5, 3.0, 2.0], 1e-5);
5292 }
5293
5294 #[test]
5295 fn minmax_value_and_active_operand() {
5296 let e = Expr::Sum(vec![
5299 Expr::MinList(vec![var(0), var(1)]),
5300 Expr::MaxList(vec![var(0), var(1)]),
5301 ]);
5302 let t = Tape::build(&e);
5303 let x = [1.3, -0.4];
5305 assert!((t.eval(&x) - (x[0] + x[1])).abs() < 1e-12);
5306 let mut g = vec![0.0; 2];
5307 t.gradient_seed(&x, 1.0, &mut g);
5308 assert!((g[0] - 1.0).abs() < 1e-12, "g0={}", g[0]);
5311 assert!((g[1] - 1.0).abs() < 1e-12, "g1={}", g[1]);
5312 }
5313
5314 #[test]
5315 fn hessian_division_matches_fd() {
5316 let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
5318 let t = Tape::build(&e);
5319 fd_check(&t, &[0.5, 1.2], 2, 1e-5);
5320 }
5321
5322 #[test]
5323 fn conditional_value_grad_hessian_active_branch() {
5324 let e = cond(
5328 cmp(CmpOp::Ge, var(0), cnst(1.0)),
5329 mul(var(0), var(1)),
5330 pow(var(1), cnst(2.0)),
5331 );
5332 let t = Tape::build(&e);
5333
5334 let x = [2.0, 5.0];
5336 assert!((t.eval(&x) - 10.0).abs() < 1e-12);
5337 let mut g = vec![0.0; 2];
5338 t.gradient_seed(&x, 1.0, &mut g);
5339 assert!((g[0] - 5.0).abs() < 1e-10);
5341 assert!((g[1] - 2.0).abs() < 1e-10);
5342 fd_check(&t, &x, 2, 1e-5);
5344
5345 let x2 = [0.0, 5.0];
5347 assert!((t.eval(&x2) - 25.0).abs() < 1e-12);
5348 let mut g2 = vec![0.0; 2];
5349 t.gradient_seed(&x2, 1.0, &mut g2);
5350 assert!(g2[0].abs() < 1e-10);
5351 assert!((g2[1] - 10.0).abs() < 1e-10);
5352 fd_check(&t, &x2, 2, 1e-5);
5353 }
5354
5355 #[test]
5356 fn comparison_and_logical_have_zero_derivative() {
5357 let lt = cmp(CmpOp::Lt, var(0), var(1));
5361 let and = Expr::And(
5362 Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
5363 Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
5364 );
5365 let notc = Expr::Not(Box::new(cmp(CmpOp::Eq, var(0), var(1))));
5366 let e = add(add(lt, and), notc);
5367 let t = Tape::build(&e);
5368
5369 let x = [1.0, 2.0];
5370 assert!((t.eval(&x) - 3.0).abs() < 1e-12);
5372 let mut g = vec![0.0; 2];
5373 t.gradient_seed(&x, 1.0, &mut g);
5374 assert!(g[0].abs() < 1e-12, "d/dx0 should be 0, got {}", g[0]);
5375 assert!(g[1].abs() < 1e-12, "d/dx1 should be 0, got {}", g[1]);
5376 }
5377
5378 #[test]
5379 fn logical_or_value() {
5380 let e = Expr::Or(
5382 Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
5383 Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
5384 );
5385 let t = Tape::build(&e);
5386 assert!((t.eval(&[-1.0, 3.0]) - 1.0).abs() < 1e-12);
5387 assert!((t.eval(&[-1.0, -3.0]) - 0.0).abs() < 1e-12);
5388 }
5389
5390 fn directional_matches_accumulate(tape: &Tape, x: &[f64], n: usize) {
5395 let vars = tape.variables();
5396 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
5397 let mut pairs = Vec::new();
5398 for (ai, &vi) in vars.iter().enumerate() {
5399 for &vj in &vars[..=ai] {
5400 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
5401 hess_map.entry((r, c)).or_insert_with(|| {
5402 let p = pairs.len();
5403 pairs.push((r, c));
5404 p
5405 });
5406 }
5407 }
5408 let nnz = pairs.len();
5409 let mut ad = vec![0.0; nnz];
5410 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
5411
5412 let nops = tape.ops.len();
5413 let mut vals = vec![0.0; nops];
5414 tape.forward_into(x, &mut vals);
5415 let mut dot = vec![0.0; nops];
5416 let mut adj = vec![0.0; nops];
5417 let mut adj_dot = vec![0.0; nops];
5418
5419 for &j in &vars {
5420 let mut seed = vec![0.0; n];
5421 seed[j] = 1.0;
5422 let mut col = vec![0.0; n];
5423 tape.hessian_directional(
5424 &vals,
5425 &seed,
5426 1.0,
5427 &mut col,
5428 &mut dot,
5429 &mut adj,
5430 &mut adj_dot,
5431 );
5432 for &i in &vars {
5433 let (r, c) = if i >= j { (i, j) } else { (j, i) };
5434 let expect = ad[hess_map[&(r, c)]];
5435 assert!(
5436 (col[i] - expect).abs() < 1e-10,
5437 "directional H[{i},{j}] = {} vs accumulate {}",
5438 col[i],
5439 expect
5440 );
5441 }
5442 }
5443 }
5444
5445 #[test]
5446 fn directional_quadratic_matches_accumulate() {
5447 let e = add(
5449 add(
5450 mul(cnst(3.0), pow(var(0), cnst(2.0))),
5451 mul(mul(cnst(2.0), var(0)), var(1)),
5452 ),
5453 pow(var(1), cnst(2.0)),
5454 );
5455 let t = Tape::build(&e);
5456 directional_matches_accumulate(&t, &[0.5, -0.3], 2);
5457 }
5458
5459 #[test]
5460 fn directional_transcendental_matches_accumulate() {
5461 let e = Expr::Sum(vec![
5462 unary(UnaryOp::Exp, var(0)),
5463 unary(UnaryOp::Sin, var(1)),
5464 unary(UnaryOp::Log, var(0)),
5465 unary(UnaryOp::Sqrt, var(1)),
5466 mul(var(0), var(1)),
5467 ]);
5468 let t = Tape::build(&e);
5469 directional_matches_accumulate(&t, &[1.5, 2.0], 2);
5470 }
5471
5472 #[test]
5473 fn directional_with_division_matches_accumulate() {
5474 let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
5475 let t = Tape::build(&e);
5476 directional_matches_accumulate(&t, &[0.5, 1.2], 2);
5477 }
5478
5479 #[test]
5480 fn hessian_sparsity_separable() {
5481 let e = add(unary(UnaryOp::Sin, var(0)), mul(var(1), var(2)));
5483 let t = Tape::build(&e);
5484 let s = t.hessian_sparsity();
5485 assert!(s.contains(&(0, 0)));
5486 assert!(s.contains(&(2, 1)));
5487 assert!(!s.contains(&(1, 0)));
5488 assert!(!s.contains(&(2, 0)));
5489 }
5490
5491 fn count_op<F: Fn(&TapeOp) -> bool>(t: &Tape, pred: F) -> usize {
5492 t.ops.iter().filter(|o| pred(o)).count()
5493 }
5494
5495 #[test]
5496 fn pow_zero_const_folds_to_one() {
5497 let e = pow(var(0), cnst(0.0));
5499 let t = Tape::build(&e);
5500 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5501 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Var(_))), 0);
5502 assert!((t.eval(&[7.0]) - 1.0).abs() < 1e-12);
5503 }
5504
5505 #[test]
5506 fn pow_one_passes_through() {
5507 let e = pow(var(0), cnst(1.0));
5509 let t = Tape::build(&e);
5510 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5511 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Const(_))), 0);
5512 assert!((t.eval(&[3.5]) - 3.5).abs() < 1e-12);
5513 }
5514
5515 #[test]
5516 fn pow_half_lowers_to_sqrt() {
5517 let e = pow(var(0), cnst(0.5));
5518 let t = Tape::build(&e);
5519 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5520 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Sqrt(_))), 1);
5521 assert!((t.eval(&[16.0]) - 4.0).abs() < 1e-12);
5522 }
5523
5524 #[test]
5525 fn pow_two_lowers_to_single_mul() {
5526 let e = pow(var(0), cnst(2.0));
5527 let t = Tape::build(&e);
5528 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5529 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5530 assert!((t.eval(&[3.0]) - 9.0).abs() < 1e-12);
5531 }
5532
5533 #[test]
5534 fn pow_three_lowers_to_two_muls() {
5535 let e = pow(var(0), cnst(3.0));
5536 let t = Tape::build(&e);
5537 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5538 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 2);
5539 assert!((t.eval(&[2.0]) - 8.0).abs() < 1e-12);
5540 }
5541
5542 #[test]
5543 fn pow_eight_lowers_to_three_muls() {
5544 let e = pow(var(0), cnst(8.0));
5546 let t = Tape::build(&e);
5547 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5548 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 3);
5549 assert!((t.eval(&[2.0]) - 256.0).abs() < 1e-12);
5550 }
5551
5552 #[test]
5553 fn pow_negative_two_lowers_to_div() {
5554 let e = pow(var(0), cnst(-2.0));
5556 let t = Tape::build(&e);
5557 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5558 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Div(..))), 1);
5559 assert!((t.eval(&[4.0]) - (1.0 / 16.0)).abs() < 1e-12);
5560 }
5561
5562 #[test]
5563 fn pow_large_const_stays_generic() {
5564 let e = pow(var(0), cnst(9.0));
5566 let t = Tape::build(&e);
5567 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5568 }
5569
5570 #[test]
5571 fn pow_non_integer_const_stays_generic() {
5572 let e = pow(var(0), cnst(1.5));
5574 let t = Tape::build(&e);
5575 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5576 }
5577
5578 #[test]
5579 fn pow_const_through_cse_const() {
5580 let two = Arc::new(cnst(2.0));
5582 let e = Expr::Binary(BinOp::Pow, Box::new(var(0)), Box::new(Expr::Cse(two)));
5583 let t = Tape::build(&e);
5584 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5585 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5586 }
5587
5588 #[test]
5589 fn hessian_pow_three_matches_fd() {
5590 let e = add(mul(cnst(5.0), pow(var(0), cnst(3.0))), mul(var(0), var(1)));
5592 let t = Tape::build(&e);
5593 fd_check(&t, &[1.7, 0.8], 2, 1e-5);
5594 }
5595
5596 #[test]
5597 fn hessian_pow_negative_matches_fd() {
5598 let e = add(pow(var(0), cnst(-2.0)), pow(var(1), cnst(2.0)));
5600 let t = Tape::build(&e);
5601 fd_check(&t, &[1.3, 2.4], 2, 1e-5);
5602 }
5603
5604 #[test]
5605 fn hessian_pow_half_matches_fd() {
5606 let e = add(pow(var(0), cnst(0.5)), mul(var(0), var(1)));
5608 let t = Tape::build(&e);
5609 fd_check(&t, &[2.5, 1.1], 2, 1e-5);
5610 }
5611
5612 #[test]
5613 fn hessian_sparsity_through_cse() {
5614 let body = Arc::new(add(var(0), var(1)));
5617 let e = add(
5618 pow(Expr::Cse(body.clone()), cnst(2.0)),
5619 Expr::Cse(body.clone()),
5620 );
5621 let t = Tape::build(&e);
5622 let s = t.hessian_sparsity();
5623 assert!(s.contains(&(0, 0)));
5624 assert!(s.contains(&(1, 0)));
5625 assert!(s.contains(&(1, 1)));
5626 assert_eq!(s.len(), 3);
5627 }
5628
5629 #[test]
5630 fn pow_forward_tangent_matches_reverse_gradient_at_base_zero() {
5631 let e = pow(var(0), var(1));
5640 let t = Tape::build(&e);
5641 assert!(
5644 t.ops.iter().any(|op| matches!(op, TapeOp::Pow(_, _))),
5645 "expected a Pow op in the tape; got {:?}",
5646 t.ops
5647 );
5648 let x = [0.0, 1.0];
5649 let n = t.ops.len();
5650
5651 let mut grad = vec![0.0; 2];
5653 t.gradient_seed(&x, 1.0, &mut grad);
5654
5655 let vals = t.forward(&x);
5657 let mut dot = vec![0.0; n];
5658 t.forward_tangent(&vals, 0, &mut dot);
5659 let fwd_dfx0 = dot[n - 1];
5660
5661 assert!(
5662 (grad[0] - 1.0).abs() < 1e-12,
5663 "reverse gradient df/dx0 at base 0 should be 1, got {}",
5664 grad[0]
5665 );
5666 assert!(
5667 (fwd_dfx0 - grad[0]).abs() < 1e-12,
5668 "forward tangent df/dx0 = {fwd_dfx0} must match reverse gradient {} at base 0",
5669 grad[0]
5670 );
5671 }
5672
5673 #[test]
5674 #[should_panic(expected = "external function calls are not supported on the")]
5675 fn hybrid_promoted_cse_with_funcall_reports_clear_message() {
5676 let body = Arc::new(Expr::Funcall {
5684 id: 0,
5685 args: vec![FuncallArg::Real(var(0))],
5686 });
5687 let exprs = vec![
5688 add(Expr::Cse(body.clone()), cnst(1.0)),
5689 add(Expr::Cse(body.clone()), cnst(2.0)),
5690 ];
5691 HybridTape::build_multi(&exprs);
5692 }
5693
5694 fn shooting_chain_roots(k: usize, leaf: Expr) -> Vec<Expr> {
5699 let step = |e: &Expr, i: usize| {
5700 let sq = mul(e.clone(), e.clone());
5701 mul(cnst(0.1), add(var(i), mul(cnst(-1.0), sq)))
5702 };
5703 let mut e = leaf;
5704 for i in 0..k - 1 {
5705 e = Expr::Cse(Arc::new(add(e.clone(), step(&e, i))));
5706 }
5707 let last = step(&e, k - 1);
5708 vec![e, last]
5709 }
5710
5711 #[test]
5716 fn hybrid_build_of_a_shared_square_chain_is_linear() {
5717 let k = 64;
5718 let roots = shooting_chain_roots(k, cnst(0.5));
5719 let tape = HybridTape::build_multi(&roots);
5720 let x: Vec<f64> = (0..k).map(|i| 0.01 * i as f64).collect();
5722 let flat = Tape::build(&add(roots[0].clone(), roots[1].clone())).eval(&x);
5723 let mut prelude = vec![0.0; tape.n_prelude_ops()];
5724 tape.forward_prelude(&x, &mut prelude);
5725 let mut local = vec![0.0; tape.max_summand_ops()];
5726 let mut hybrid = 0.0;
5727 for s in &tape.summands {
5728 tape.forward_summand(s, &x, &prelude, &mut local);
5729 hybrid += tape.root_value(s, &local);
5730 }
5731 assert!(
5732 (hybrid - flat).abs() <= 1e-12 * flat.abs().max(1.0),
5733 "{hybrid} vs {flat}"
5734 );
5735 }
5736
5737 #[test]
5741 #[should_panic(expected = "external function calls are not supported on the")]
5742 fn hybrid_funcall_under_a_shared_square_chain_still_reports() {
5743 let leaf = Expr::Funcall {
5744 id: 0,
5745 args: vec![FuncallArg::Real(var(0))],
5746 };
5747 HybridTape::build_multi(&shooting_chain_roots(64, leaf));
5748 }
5749
5750 #[test]
5761 fn directional_hybrid_hessian_matches_flat_directional() {
5762 let body = Arc::new(unary(UnaryOp::Exp, mul(cnst(0.5), add(var(0), var(1)))));
5765 let exprs = vec![
5766 pow(Expr::Cse(body.clone()), cnst(2.0)),
5767 mul(Expr::Cse(body.clone()), var(2)),
5768 add(
5769 unary(UnaryOp::Sin, Expr::Cse(body.clone())),
5770 pow(var(2), cnst(2.0)),
5771 ),
5772 ];
5773 let weights = [1.25, -0.75, 2.5];
5774 let x = [0.3, -0.1, 0.7];
5775 let n = 3;
5776
5777 let hybrid = HybridTape::build_multi(&exprs);
5778 assert!(
5779 hybrid.n_prelude_ops() > 0,
5780 "a CSE shared by 3 roots must be promoted into the prelude"
5781 );
5782
5783 let seeds = [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]];
5785 for seed in &seeds {
5786 let mut flat_out = vec![0.0; n];
5787 for (e, &wt) in exprs.iter().zip(&weights) {
5788 let t = Tape::build(e);
5789 let vals = t.forward(&x);
5790 let m = t.ops.len();
5791 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; m], vec![0.0; m], vec![0.0; m]);
5792 t.hessian_directional(
5793 &vals,
5794 seed,
5795 wt,
5796 &mut flat_out,
5797 &mut dot,
5798 &mut adj,
5799 &mut adj_dot,
5800 );
5801 }
5802
5803 let mut hyb_out = vec![0.0; n];
5804 let np = hybrid.n_prelude_ops();
5805 let ml = hybrid.max_summand_ops();
5806 let mut prelude_vals = vec![0.0; np];
5807 let mut prelude_dot = vec![0.0; np];
5808 let mut prelude_adj = vec![0.0; np];
5809 let mut prelude_adj_dot = vec![0.0; np];
5810 let (mut local_dot, mut local_adj, mut local_adj_dot) =
5811 (vec![0.0; ml], vec![0.0; ml], vec![0.0; ml]);
5812 let creach: Vec<u32> = {
5818 let mut u: BTreeSet<u32> = BTreeSet::new();
5819 for s in &hybrid.summands {
5820 u.extend(s.prelude_reach.iter().map(|&p| p as u32));
5821 }
5822 u.into_iter().collect()
5823 };
5824 hybrid.forward_prelude(&x, &mut prelude_vals);
5825 hybrid.prelude_tangent(&prelude_vals, seed, &creach, &mut prelude_dot);
5826 for (s, &wt) in hybrid.summands.iter().zip(&weights) {
5827 let mut local_vals = vec![0.0; s.ops.len()];
5828 hybrid.forward_summand(s, &x, &prelude_vals, &mut local_vals);
5829 hybrid.hessian_summand_directional(
5830 s,
5831 &local_vals,
5832 &prelude_dot,
5833 seed,
5834 wt,
5835 &mut hyb_out,
5836 &mut local_dot,
5837 &mut local_adj,
5838 &mut local_adj_dot,
5839 &mut prelude_adj,
5840 &mut prelude_adj_dot,
5841 );
5842 }
5843 hybrid.prelude_reverse_directional(
5844 &prelude_vals,
5845 &prelude_dot,
5846 &creach,
5847 &mut hyb_out,
5848 &mut prelude_adj,
5849 &mut prelude_adj_dot,
5850 );
5851
5852 for k in 0..n {
5853 let scale = flat_out[k].abs().max(1.0);
5854 assert!(
5855 (hyb_out[k] - flat_out[k]).abs() <= 1e-13 * scale,
5856 "seed {seed:?} entry {k}: hybrid {} vs flat {}",
5857 hyb_out[k],
5858 flat_out[k]
5859 );
5860 }
5861 assert!(
5862 hyb_out.iter().any(|v| *v != 0.0),
5863 "an all-zero H·s is no comparison"
5864 );
5865 assert!(prelude_adj.iter().all(|v| *v == 0.0));
5868 assert!(prelude_adj_dot.iter().all(|v| *v == 0.0));
5869 }
5870 }
5871}