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]
256pub(crate) fn erf(u: f64) -> f64 {
257 libm::erf(u)
258}
259
260#[inline]
265pub(crate) fn erf_d1(u: f64) -> f64 {
266 std::f64::consts::FRAC_2_SQRT_PI * (-u * u).exp()
267}
268
269#[inline]
280pub(crate) fn erf_d2(u: f64) -> f64 {
281 -2.0 * (u * erf_d1(u))
282}
283
284#[inline]
287fn cmp_holds(op: CmpOp, a: f64, b: f64) -> bool {
288 match op {
289 CmpOp::Lt => a < b,
290 CmpOp::Le => a <= b,
291 CmpOp::Eq => a == b,
292 CmpOp::Ge => a >= b,
293 CmpOp::Gt => a > b,
294 CmpOp::Ne => a != b,
295 }
296}
297
298fn funcall_to_ext_args<'a>(args: &'a [TapeFuncallArg], vals: &[f64]) -> Vec<ExternalArg<'a>> {
299 args.iter()
300 .map(|a| match a {
301 TapeFuncallArg::Tape(idx) => ExternalArg::Real(vals[*idx]),
302 TapeFuncallArg::Str(s) => ExternalArg::Str(s.as_str()),
303 })
304 .collect()
305}
306
307fn ext_eval_or_nan(
320 lib: &ExternalLibrary,
321 name: &str,
322 call_args: &[ExternalArg<'_>],
323 n_args: usize,
324 want_derivs: bool,
325 want_hes: bool,
326) -> EvalResult {
327 lib.eval(name, call_args, want_derivs, want_hes)
328 .unwrap_or_else(|_| EvalResult {
329 value: f64::NAN,
330 derivs: want_derivs.then(|| vec![f64::NAN; n_args]),
331 hessian: want_hes.then(|| vec![f64::NAN; n_args * (n_args + 1) / 2]),
332 })
333}
334
335#[derive(Debug, Clone)]
338pub struct Tape {
339 pub ops: Vec<TapeOp>,
340}
341
342impl Tape {
343 pub fn build(expr: &Expr) -> Self {
347 Self::build_with_externals(expr, &ExternalResolver::default())
348 }
349
350 pub fn build_with_externals(expr: &Expr, resolver: &ExternalResolver) -> Self {
355 let mut ops = Vec::new();
356 let mut cache: HashMap<*const Expr, usize> = HashMap::new();
357 build_recursive(expr, &mut ops, &mut cache, resolver);
358 Tape { ops }
359 }
360
361 pub fn forward(&self, x: &[f64]) -> Vec<f64> {
364 let mut vals: Vec<f64> = Vec::with_capacity(self.ops.len());
365 for op in &self.ops {
366 let v = match op {
367 TapeOp::Const(c) => *c,
368 TapeOp::Var(i) => x[*i],
369 TapeOp::Add(a, b) => vals[*a] + vals[*b],
370 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
371 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
372 TapeOp::Div(a, b) => vals[*a] / vals[*b],
373 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
374 TapeOp::Neg(a) => -vals[*a],
375 TapeOp::Abs(a) => vals[*a].abs(),
376 TapeOp::Sqrt(a) => vals[*a].sqrt(),
377 TapeOp::Exp(a) => vals[*a].exp(),
378 TapeOp::Log(a) => vals[*a].ln(),
379 TapeOp::Log10(a) => vals[*a].log10(),
380 TapeOp::Sin(a) => vals[*a].sin(),
381 TapeOp::Cos(a) => vals[*a].cos(),
382 TapeOp::Tan(a) => vals[*a].tan(),
383 TapeOp::Atan(a) => vals[*a].atan(),
384 TapeOp::Acos(a) => vals[*a].acos(),
385 TapeOp::Sinh(a) => vals[*a].sinh(),
386 TapeOp::Cosh(a) => vals[*a].cosh(),
387 TapeOp::Tanh(a) => vals[*a].tanh(),
388 TapeOp::Asin(a) => vals[*a].asin(),
389 TapeOp::Acosh(a) => vals[*a].acosh(),
390 TapeOp::Asinh(a) => vals[*a].asinh(),
391 TapeOp::Atanh(a) => vals[*a].atanh(),
392 TapeOp::Erf(a) => erf(vals[*a]),
393 TapeOp::XLogX(a) => xlogx(vals[*a]),
394 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
395 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
396 TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
397 TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
398 TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
399 TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
400 TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
401 TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
402 TapeOp::Select(c, t, e) => {
403 if vals[*c] != 0.0 {
404 vals[*t]
405 } else {
406 vals[*e]
407 }
408 }
409 TapeOp::Funcall(fc) => {
410 let FuncallData { lib, name, args } = fc.as_ref();
411 let call_args = funcall_to_ext_args(args, &vals);
412 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
413 res.value
414 }
415 };
416 vals.push(v);
417 }
418 vals
419 }
420
421 pub fn eval(&self, x: &[f64]) -> f64 {
422 let vals = self.forward(x);
423 *vals.last().unwrap_or(&0.0)
424 }
425
426 pub fn gradient_seed(&self, x: &[f64], seed: f64, grad: &mut [f64]) {
431 if seed == 0.0 || self.ops.is_empty() {
432 return;
433 }
434 let vals = self.forward(x);
435 self.reverse(&vals, seed, grad);
436 }
437
438 pub fn gradient_seed_into(
450 &self,
451 x: &[f64],
452 seed: f64,
453 grad: &mut [f64],
454 vals: &mut [f64],
455 adj: &mut [f64],
456 ) {
457 if seed == 0.0 || self.ops.is_empty() {
458 return;
459 }
460 debug_assert!(vals.len() >= self.ops.len());
461 self.forward_into(x, vals);
462 self.reverse_into(vals, seed, grad, adj);
463 }
464
465 fn reverse(&self, vals: &[f64], seed: f64, grad: &mut [f64]) {
466 let n = self.ops.len();
467 let mut adj = vec![0.0f64; n];
468 self.reverse_into(vals, seed, grad, &mut adj);
469 }
470
471 fn reverse_into(&self, vals: &[f64], seed: f64, grad: &mut [f64], adj: &mut [f64]) {
476 let n = self.ops.len();
477 debug_assert!(adj.len() >= n);
478 adj[..n].fill(0.0);
479 adj[n - 1] = seed;
480
481 for i in (0..n).rev() {
482 let a = adj[i];
483 if a == 0.0 {
484 continue;
485 }
486 match &self.ops[i] {
487 TapeOp::Const(_) => {}
488 TapeOp::Var(j) => {
489 grad[*j] += a;
490 }
491 TapeOp::Add(l, r) => {
492 adj[*l] += a;
493 adj[*r] += a;
494 }
495 TapeOp::Sub(l, r) => {
496 adj[*l] += a;
497 adj[*r] -= a;
498 }
499 TapeOp::Mul(l, r) => {
500 adj[*l] += a * vals[*r];
501 adj[*r] += a * vals[*l];
502 }
503 TapeOp::Div(l, r) => {
504 let rv = vals[*r];
508 adj[*l] += a / rv;
509 adj[*r] -= a * vals[i] / rv;
510 }
511 TapeOp::Pow(l, r) => {
512 let lv = vals[*l];
513 let rv = vals[*r];
514 if rv != 0.0 {
515 adj[*l] += a * rv * lv.powf(rv - 1.0);
516 }
517 if lv > 0.0 {
518 adj[*r] += a * vals[i] * lv.ln();
519 }
520 }
521 TapeOp::Neg(j) => {
522 adj[*j] -= a;
523 }
524 TapeOp::Abs(j) => {
525 if vals[*j] >= 0.0 {
526 adj[*j] += a;
527 } else {
528 adj[*j] -= a;
529 }
530 }
531 TapeOp::Sqrt(j) => {
532 let sv = vals[i];
533 if sv > 0.0 {
534 adj[*j] += a * 0.5 / sv;
535 }
536 }
537 TapeOp::Exp(j) => {
538 adj[*j] += a * vals[i];
539 }
540 TapeOp::Log(j) => {
541 adj[*j] += a / vals[*j];
542 }
543 TapeOp::Log10(j) => {
544 adj[*j] += a / (vals[*j] * std::f64::consts::LN_10);
545 }
546 TapeOp::Sin(j) => {
547 adj[*j] += a * vals[*j].cos();
548 }
549 TapeOp::Cos(j) => {
550 adj[*j] -= a * vals[*j].sin();
551 }
552 TapeOp::Tan(j) => {
553 let t = vals[i];
554 adj[*j] += a * (1.0 + t * t);
555 }
556 TapeOp::Atan(j) => {
557 let u = vals[*j];
558 adj[*j] += a / (1.0 + u * u);
559 }
560 TapeOp::Acos(j) => {
561 let u = vals[*j];
562 adj[*j] -= a / (1.0 - u * u).sqrt();
563 }
564 TapeOp::Sinh(j) => {
565 adj[*j] += a * vals[*j].cosh();
566 }
567 TapeOp::Cosh(j) => {
568 adj[*j] += a * vals[*j].sinh();
569 }
570 TapeOp::Tanh(j) => {
571 let t = vals[i];
572 adj[*j] += a * (1.0 - t * t);
573 }
574 TapeOp::Asin(j) => {
575 let u = vals[*j];
576 adj[*j] += a / (1.0 - u * u).sqrt();
577 }
578 TapeOp::Acosh(j) => {
579 let u = vals[*j];
580 adj[*j] += a / (u * u - 1.0).sqrt();
581 }
582 TapeOp::Asinh(j) => {
583 let u = vals[*j];
584 adj[*j] += a / (u * u + 1.0).sqrt();
585 }
586 TapeOp::Atanh(j) => {
587 let u = vals[*j];
588 adj[*j] += a / (1.0 - u * u);
589 }
590 TapeOp::Erf(j) => {
591 adj[*j] += a * erf_d1(vals[*j]);
592 }
593 TapeOp::XLogX(j) => {
594 adj[*j] += a * xlogx_d1(vals[*j]);
595 }
596 TapeOp::CEntropy(l, r) => {
597 adj[*l] += a * centropy_da(vals[*l], vals[*r]);
598 adj[*r] += a * centropy_db(vals[*l], vals[*r]);
599 }
600 TapeOp::Atan2(l, r) => {
601 let y = vals[*l];
602 let x = vals[*r];
603 let d = y * y + x * x;
604 adj[*l] += a * (x / d);
605 adj[*r] += a * (-y / d);
606 }
607 TapeOp::Min(l, r) => {
611 if vals[*l] <= vals[*r] {
612 adj[*l] += a;
613 } else {
614 adj[*r] += a;
615 }
616 }
617 TapeOp::Max(l, r) => {
618 if vals[*l] >= vals[*r] {
619 adj[*l] += a;
620 } else {
621 adj[*r] += a;
622 }
623 }
624 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
627 TapeOp::Select(c, t, e) => {
630 if vals[*c] != 0.0 {
631 adj[*t] += a;
632 } else {
633 adj[*e] += a;
634 }
635 }
636 TapeOp::Funcall(fc) => {
637 let FuncallData { lib, name, args } = fc.as_ref();
638 let call_args = funcall_to_ext_args(args, vals);
639 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
640 let derivs = res.derivs.expect("want_derivs=true returns derivs");
641 let mut k = 0usize;
642 for arg in args {
643 if let TapeFuncallArg::Tape(idx) = arg {
644 adj[*idx] += a * derivs[k];
645 k += 1;
646 }
647 }
648 }
649 }
650 }
651 }
652
653 pub fn variables(&self) -> Vec<usize> {
655 let mut s: BTreeSet<usize> = BTreeSet::new();
656 for op in &self.ops {
657 if let TapeOp::Var(j) = op {
658 s.insert(*j);
659 }
660 }
661 s.into_iter().collect()
662 }
663
664 fn forward_tangent(&self, vals: &[f64], seed_var: usize, dot: &mut [f64]) {
669 let n = self.ops.len();
670 debug_assert_eq!(dot.len(), n);
671 for i in 0..n {
672 dot[i] = match &self.ops[i] {
673 TapeOp::Const(_) => 0.0,
674 TapeOp::Var(k) => {
675 if *k == seed_var {
676 1.0
677 } else {
678 0.0
679 }
680 }
681 TapeOp::Add(a, b) => dot[*a] + dot[*b],
682 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
683 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
684 TapeOp::Div(a, b) => {
685 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
693 }
694 TapeOp::Pow(a, b) => {
695 let u = vals[*a];
696 let r = vals[*b];
697 let du = dot[*a];
698 let dr = dot[*b];
699 let mut result = 0.0;
700 if r != 0.0 {
705 result += r * u.powf(r - 1.0) * du;
706 }
707 if u > 0.0 {
708 result += vals[i] * u.ln() * dr;
709 }
710 result
711 }
712 TapeOp::Neg(a) => -dot[*a],
713 TapeOp::Abs(a) => {
714 if vals[*a] >= 0.0 {
715 dot[*a]
716 } else {
717 -dot[*a]
718 }
719 }
720 TapeOp::Sqrt(a) => {
721 let sv = vals[i];
722 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
723 }
724 TapeOp::Exp(a) => dot[*a] * vals[i],
725 TapeOp::Log(a) => dot[*a] / vals[*a],
726 TapeOp::Log10(a) => dot[*a] / (vals[*a] * std::f64::consts::LN_10),
727 TapeOp::Sin(a) => dot[*a] * vals[*a].cos(),
728 TapeOp::Cos(a) => -dot[*a] * vals[*a].sin(),
729 TapeOp::Tan(a) => {
730 let t = vals[i];
731 dot[*a] * (1.0 + t * t)
732 }
733 TapeOp::Atan(a) => {
734 let u = vals[*a];
735 dot[*a] / (1.0 + u * u)
736 }
737 TapeOp::Acos(a) => {
738 let u = vals[*a];
739 -dot[*a] / (1.0 - u * u).sqrt()
740 }
741 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
742 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
743 TapeOp::Tanh(a) => {
744 let t = vals[i];
745 dot[*a] * (1.0 - t * t)
746 }
747 TapeOp::Asin(a) => {
748 let u = vals[*a];
749 dot[*a] / (1.0 - u * u).sqrt()
750 }
751 TapeOp::Acosh(a) => {
752 let u = vals[*a];
753 dot[*a] / (u * u - 1.0).sqrt()
754 }
755 TapeOp::Asinh(a) => {
756 let u = vals[*a];
757 dot[*a] / (u * u + 1.0).sqrt()
758 }
759 TapeOp::Atanh(a) => {
760 let u = vals[*a];
761 dot[*a] / (1.0 - u * u)
762 }
763 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
764 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
765 TapeOp::CEntropy(a, b) => {
766 centropy_da(vals[*a], vals[*b]) * dot[*a]
767 + centropy_db(vals[*a], vals[*b]) * dot[*b]
768 }
769 TapeOp::Atan2(a, b) => {
770 let y = vals[*a];
771 let x = vals[*b];
772 let d = y * y + x * x;
773 (x * dot[*a] - y * dot[*b]) / d
774 }
775 TapeOp::Min(a, b) => {
777 if vals[*a] <= vals[*b] {
778 dot[*a]
779 } else {
780 dot[*b]
781 }
782 }
783 TapeOp::Max(a, b) => {
784 if vals[*a] >= vals[*b] {
785 dot[*a]
786 } else {
787 dot[*b]
788 }
789 }
790 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
791 TapeOp::Select(c, t, e) => {
792 if vals[*c] != 0.0 {
793 dot[*t]
794 } else {
795 dot[*e]
796 }
797 }
798 TapeOp::Funcall(fc) => {
799 let FuncallData { lib, name, args } = fc.as_ref();
800 let call_args = funcall_to_ext_args(args, vals);
801 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
802 let derivs = res.derivs.expect("want_derivs=true returns derivs");
803 let mut acc = 0.0;
804 let mut k = 0usize;
805 for arg in args {
806 if let TapeFuncallArg::Tape(idx) = arg {
807 acc += derivs[k] * dot[*idx];
808 k += 1;
809 }
810 }
811 acc
812 }
813 };
814 }
815 }
816
817 pub fn forward_into(&self, x: &[f64], vals: &mut [f64]) {
821 let n = self.ops.len();
822 debug_assert!(vals.len() >= n);
823 for i in 0..n {
824 vals[i] = match &self.ops[i] {
825 TapeOp::Const(c) => *c,
826 TapeOp::Var(j) => x[*j],
827 TapeOp::Add(a, b) => vals[*a] + vals[*b],
828 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
829 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
830 TapeOp::Div(a, b) => vals[*a] / vals[*b],
831 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
832 TapeOp::Neg(a) => -vals[*a],
833 TapeOp::Abs(a) => vals[*a].abs(),
834 TapeOp::Sqrt(a) => vals[*a].sqrt(),
835 TapeOp::Exp(a) => vals[*a].exp(),
836 TapeOp::Log(a) => vals[*a].ln(),
837 TapeOp::Log10(a) => vals[*a].log10(),
838 TapeOp::Sin(a) => vals[*a].sin(),
839 TapeOp::Cos(a) => vals[*a].cos(),
840 TapeOp::Tan(a) => vals[*a].tan(),
841 TapeOp::Atan(a) => vals[*a].atan(),
842 TapeOp::Acos(a) => vals[*a].acos(),
843 TapeOp::Sinh(a) => vals[*a].sinh(),
844 TapeOp::Cosh(a) => vals[*a].cosh(),
845 TapeOp::Tanh(a) => vals[*a].tanh(),
846 TapeOp::Asin(a) => vals[*a].asin(),
847 TapeOp::Acosh(a) => vals[*a].acosh(),
848 TapeOp::Asinh(a) => vals[*a].asinh(),
849 TapeOp::Atanh(a) => vals[*a].atanh(),
850 TapeOp::Erf(a) => erf(vals[*a]),
851 TapeOp::XLogX(a) => xlogx(vals[*a]),
852 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
853 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
854 TapeOp::Min(a, b) => vals[*a].min(vals[*b]),
855 TapeOp::Max(a, b) => vals[*a].max(vals[*b]),
856 TapeOp::Cmp(op, a, b) => f64::from(cmp_holds(*op, vals[*a], vals[*b])),
857 TapeOp::And(a, b) => f64::from(vals[*a] != 0.0 && vals[*b] != 0.0),
858 TapeOp::Or(a, b) => f64::from(vals[*a] != 0.0 || vals[*b] != 0.0),
859 TapeOp::Not(a) => f64::from(vals[*a] == 0.0),
860 TapeOp::Select(c, t, e) => {
861 if vals[*c] != 0.0 {
862 vals[*t]
863 } else {
864 vals[*e]
865 }
866 }
867 TapeOp::Funcall(fc) => {
868 let FuncallData { lib, name, args } = fc.as_ref();
869 let call_args = funcall_to_ext_args(args, &*vals);
870 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), false, false);
871 res.value
872 }
873 };
874 }
875 }
876
877 pub fn eval_into(&self, x: &[f64], vals: &mut [f64]) -> f64 {
888 let n = self.ops.len();
889 if n == 0 {
890 return 0.0;
891 }
892 self.forward_into(x, vals);
893 vals[n - 1]
894 }
895
896 pub fn hessian_directional(
913 &self,
914 vals: &[f64],
915 seed: &[f64],
916 weight: f64,
917 out: &mut [f64],
918 dot: &mut [f64],
919 adj: &mut [f64],
920 adj_dot: &mut [f64],
921 ) {
922 let n = self.ops.len();
923 if n == 0 || weight == 0.0 {
924 return;
925 }
926 debug_assert!(vals.len() >= n);
927 debug_assert!(dot.len() >= n);
928 debug_assert!(adj.len() >= n);
929 debug_assert!(adj_dot.len() >= n);
930
931 for i in 0..n {
935 dot[i] = match &self.ops[i] {
936 TapeOp::Const(_) => 0.0,
937 TapeOp::Var(k) => seed[*k],
938 TapeOp::Add(a, b) => dot[*a] + dot[*b],
939 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
940 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
941 TapeOp::Div(a, b) => {
942 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
950 }
951 TapeOp::Pow(a, b) => {
952 let u = vals[*a];
953 let r = vals[*b];
954 let du = dot[*a];
955 let dr = dot[*b];
956 let mut result = 0.0;
957 if r != 0.0 {
962 result += r * u.powf(r - 1.0) * du;
963 }
964 if u > 0.0 {
965 result += vals[i] * u.ln() * dr;
966 }
967 result
968 }
969 TapeOp::Neg(a) => -dot[*a],
970 TapeOp::Abs(a) => {
971 if vals[*a] >= 0.0 {
972 dot[*a]
973 } else {
974 -dot[*a]
975 }
976 }
977 TapeOp::Sqrt(a) => {
978 let sv = vals[i];
979 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
980 }
981 TapeOp::Exp(a) => vals[i] * dot[*a],
982 TapeOp::Log(a) => dot[*a] / vals[*a],
983 TapeOp::Log10(a) => dot[*a] / (vals[*a] * std::f64::consts::LN_10),
984 TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
985 TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
986 TapeOp::Tan(a) => {
987 let t = vals[i];
988 (1.0 + t * t) * dot[*a]
989 }
990 TapeOp::Atan(a) => {
991 let u = vals[*a];
992 dot[*a] / (1.0 + u * u)
993 }
994 TapeOp::Acos(a) => {
995 let u = vals[*a];
996 -dot[*a] / (1.0 - u * u).sqrt()
997 }
998 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
999 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
1000 TapeOp::Tanh(a) => {
1001 let t = vals[i];
1002 (1.0 - t * t) * dot[*a]
1003 }
1004 TapeOp::Asin(a) => {
1005 let u = vals[*a];
1006 dot[*a] / (1.0 - u * u).sqrt()
1007 }
1008 TapeOp::Acosh(a) => {
1009 let u = vals[*a];
1010 dot[*a] / (u * u - 1.0).sqrt()
1011 }
1012 TapeOp::Asinh(a) => {
1013 let u = vals[*a];
1014 dot[*a] / (u * u + 1.0).sqrt()
1015 }
1016 TapeOp::Atanh(a) => {
1017 let u = vals[*a];
1018 dot[*a] / (1.0 - u * u)
1019 }
1020 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
1021 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
1022 TapeOp::CEntropy(a, b) => {
1023 centropy_da(vals[*a], vals[*b]) * dot[*a]
1024 + centropy_db(vals[*a], vals[*b]) * dot[*b]
1025 }
1026 TapeOp::Atan2(a, b) => {
1027 let y = vals[*a];
1028 let x = vals[*b];
1029 let d = y * y + x * x;
1030 (x * dot[*a] - y * dot[*b]) / d
1031 }
1032 TapeOp::Min(a, b) => {
1034 if vals[*a] <= vals[*b] {
1035 dot[*a]
1036 } else {
1037 dot[*b]
1038 }
1039 }
1040 TapeOp::Max(a, b) => {
1041 if vals[*a] >= vals[*b] {
1042 dot[*a]
1043 } else {
1044 dot[*b]
1045 }
1046 }
1047 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => 0.0,
1048 TapeOp::Select(c, t, e) => {
1049 if vals[*c] != 0.0 {
1050 dot[*t]
1051 } else {
1052 dot[*e]
1053 }
1054 }
1055 TapeOp::Funcall(fc) => {
1056 let FuncallData { lib, name, args } = fc.as_ref();
1057 let call_args = funcall_to_ext_args(args, vals);
1058 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, false);
1059 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1060 let mut acc = 0.0;
1061 let mut k = 0usize;
1062 for arg in args {
1063 if let TapeFuncallArg::Tape(idx) = arg {
1064 acc += derivs[k] * dot[*idx];
1065 k += 1;
1066 }
1067 }
1068 acc
1069 }
1070 };
1071 }
1072
1073 for slot in adj.iter_mut().take(n) {
1077 *slot = 0.0;
1078 }
1079 for slot in adj_dot.iter_mut().take(n) {
1080 *slot = 0.0;
1081 }
1082 adj[n - 1] = 1.0;
1083
1084 for i in (0..n).rev() {
1085 let w = adj[i];
1086 let wd = adj_dot[i];
1087 if w == 0.0 && wd == 0.0 {
1088 continue;
1089 }
1090 match &self.ops[i] {
1091 TapeOp::Const(_) => {}
1092 TapeOp::Var(k) => {
1093 if wd != 0.0 {
1094 out[*k] += weight * wd;
1095 }
1096 }
1097 TapeOp::Add(a, b) => {
1098 adj[*a] += w;
1099 adj[*b] += w;
1100 adj_dot[*a] += wd;
1101 adj_dot[*b] += wd;
1102 }
1103 TapeOp::Sub(a, b) => {
1104 adj[*a] += w;
1105 adj[*b] -= w;
1106 adj_dot[*a] += wd;
1107 adj_dot[*b] -= wd;
1108 }
1109 TapeOp::Mul(a, b) => {
1110 adj[*a] += w * vals[*b];
1111 adj[*b] += w * vals[*a];
1112 adj_dot[*a] += wd * vals[*b] + w * dot[*b];
1113 adj_dot[*b] += wd * vals[*a] + w * dot[*a];
1114 }
1115 TapeOp::Div(a, b) => {
1116 let vb = vals[*b];
1123 let q = vals[i];
1124 let qd = dot[i];
1125 adj[*a] += w / vb;
1126 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1127 adj[*b] -= w * q / vb;
1128 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1129 }
1130 TapeOp::Pow(a, b) => {
1131 let u = vals[*a];
1132 let r = vals[*b];
1133 let du = dot[*a];
1134 let dr = dot[*b];
1135 if r != 0.0 {
1136 if u != 0.0 {
1137 let p_a = r * u.powf(r - 1.0);
1138 adj[*a] += w * p_a;
1139 let mut dp_a = dr * u.powf(r - 1.0);
1140 if u > 0.0 {
1141 dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1142 } else {
1143 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1144 }
1145 adj_dot[*a] += wd * p_a + w * dp_a;
1146 } else if r >= 2.0 {
1147 let p_a = 0.0;
1148 adj[*a] += w * p_a;
1149 let dp_a = if r == 2.0 {
1150 2.0 * du
1151 } else {
1152 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1153 };
1154 adj_dot[*a] += wd * p_a + w * dp_a;
1155 }
1156 }
1157 if u > 0.0 {
1158 let ln_u = u.ln();
1159 let p_b = vals[i] * ln_u;
1160 adj[*b] += w * p_b;
1161 let dur = vals[i] * (r * du / u + dr * ln_u);
1162 let dp_b = dur * ln_u + vals[i] * du / u;
1163 adj_dot[*b] += wd * p_b + w * dp_b;
1164 }
1165 }
1166 TapeOp::Neg(a) => {
1167 adj[*a] -= w;
1168 adj_dot[*a] -= wd;
1169 }
1170 TapeOp::Abs(a) => {
1171 let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
1172 adj[*a] += w * s;
1173 adj_dot[*a] += wd * s;
1174 }
1175 TapeOp::Sqrt(a) => {
1176 let sv = vals[i];
1177 if sv > 0.0 {
1178 let fp = 0.5 / sv;
1179 let fpp = -0.25 / (vals[*a] * sv);
1180 adj[*a] += w * fp;
1181 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1182 }
1183 }
1184 TapeOp::Exp(a) => {
1185 let ev = vals[i];
1186 adj[*a] += w * ev;
1187 adj_dot[*a] += wd * ev + w * ev * dot[*a];
1188 }
1189 TapeOp::Log(a) => {
1190 let u = vals[*a];
1191 adj[*a] += w / u;
1192 adj_dot[*a] += wd / u + w * (-1.0 / (u * u)) * dot[*a];
1193 }
1194 TapeOp::Log10(a) => {
1195 let u = vals[*a];
1196 let c = std::f64::consts::LN_10;
1197 adj[*a] += w / (u * c);
1198 adj_dot[*a] += wd / (u * c) + w * (-1.0 / (u * u * c)) * dot[*a];
1199 }
1200 TapeOp::Sin(a) => {
1201 let u = vals[*a];
1202 let cu = u.cos();
1203 adj[*a] += w * cu;
1204 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1205 }
1206 TapeOp::Cos(a) => {
1207 let u = vals[*a];
1208 let su = u.sin();
1209 adj[*a] -= w * su;
1210 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1211 }
1212 TapeOp::Tan(a) => {
1213 let t = vals[i];
1214 let gp = 1.0 + t * t;
1215 let gpp = 2.0 * t * gp;
1216 adj[*a] += w * gp;
1217 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1218 }
1219 TapeOp::Atan(a) => {
1220 let u = vals[*a];
1221 let d = 1.0 + u * u;
1222 let gp = 1.0 / d;
1223 let gpp = -2.0 * u / (d * d);
1224 adj[*a] += w * gp;
1225 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1226 }
1227 TapeOp::Acos(a) => {
1228 let u = vals[*a];
1229 let s = 1.0 - u * u;
1230 let r = s.sqrt();
1231 let gp = -1.0 / r;
1232 let gpp = -u / (s * r);
1233 adj[*a] += w * gp;
1234 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1235 }
1236 TapeOp::Sinh(a) => {
1237 let u = vals[*a];
1238 let gp = u.cosh();
1239 let gpp = u.sinh();
1240 adj[*a] += w * gp;
1241 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1242 }
1243 TapeOp::Cosh(a) => {
1244 let u = vals[*a];
1245 let gp = u.sinh();
1246 let gpp = u.cosh();
1247 adj[*a] += w * gp;
1248 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1249 }
1250 TapeOp::Tanh(a) => {
1251 let t = vals[i];
1252 let gp = 1.0 - t * t;
1253 let gpp = -2.0 * t * gp;
1254 adj[*a] += w * gp;
1255 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1256 }
1257 TapeOp::Asin(a) => {
1258 let u = vals[*a];
1259 let s = 1.0 - u * u;
1260 let r = s.sqrt();
1261 let gp = 1.0 / r;
1262 let gpp = u / (s * r);
1263 adj[*a] += w * gp;
1264 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1265 }
1266 TapeOp::Acosh(a) => {
1267 let u = vals[*a];
1268 let s = u * u - 1.0;
1269 let r = s.sqrt();
1270 let gp = 1.0 / r;
1271 let gpp = -u / (s * r);
1272 adj[*a] += w * gp;
1273 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1274 }
1275 TapeOp::Asinh(a) => {
1276 let u = vals[*a];
1277 let s = u * u + 1.0;
1278 let r = s.sqrt();
1279 let gp = 1.0 / r;
1280 let gpp = -u / (s * r);
1281 adj[*a] += w * gp;
1282 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1283 }
1284 TapeOp::Atanh(a) => {
1285 let u = vals[*a];
1286 let d = 1.0 - u * u;
1287 let gp = 1.0 / d;
1288 let gpp = 2.0 * u / (d * d);
1289 adj[*a] += w * gp;
1290 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1291 }
1292 TapeOp::Erf(a) => {
1293 let u = vals[*a];
1294 let gp = erf_d1(u);
1295 let gpp = erf_d2(u);
1296 adj[*a] += w * gp;
1297 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1298 }
1299 TapeOp::XLogX(a) => {
1300 let u = vals[*a];
1303 let gp = xlogx_d1(u);
1304 let gpp = xlogx_d2(u);
1305 adj[*a] += w * gp;
1306 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1307 }
1308 TapeOp::CEntropy(a, b) => {
1309 let ua = vals[*a];
1310 let ub = vals[*b];
1311 let fa = centropy_da(ua, ub);
1312 let fb = centropy_db(ua, ub);
1313 let faa = centropy_daa(ua);
1314 let fab = centropy_dab(ub);
1315 let fbb = centropy_dbb(ua, ub);
1316 adj[*a] += w * fa;
1317 adj[*b] += w * fb;
1318 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1319 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1320 }
1321 TapeOp::Atan2(a, b) => {
1322 let y = vals[*a];
1323 let x = vals[*b];
1324 let d = y * y + x * x;
1325 let d2 = d * d;
1326 let fa = x / d;
1327 let fb = -y / d;
1328 let faa = -2.0 * y * x / d2;
1329 let fab = (y * y - x * x) / d2;
1330 let fbb = 2.0 * y * x / d2;
1331 adj[*a] += w * fa;
1332 adj[*b] += w * fb;
1333 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1334 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1335 }
1336 TapeOp::Min(a, b) => {
1340 let br = if vals[*a] <= vals[*b] { *a } else { *b };
1341 adj[br] += w;
1342 adj_dot[br] += wd;
1343 }
1344 TapeOp::Max(a, b) => {
1345 let br = if vals[*a] >= vals[*b] { *a } else { *b };
1346 adj[br] += w;
1347 adj_dot[br] += wd;
1348 }
1349 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {}
1351 TapeOp::Select(c, t, e) => {
1354 let br = if vals[*c] != 0.0 { *t } else { *e };
1355 adj[br] += w;
1356 adj_dot[br] += wd;
1357 }
1358 TapeOp::Funcall(fc) => {
1359 let FuncallData { lib, name, args } = fc.as_ref();
1360 let call_args = funcall_to_ext_args(args, vals);
1361 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1362 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1363 let hes = res.hessian.expect("want_hes=true returns hessian");
1364 let real_tape: Vec<usize> = args
1365 .iter()
1366 .filter_map(|a| match a {
1367 TapeFuncallArg::Tape(t) => Some(*t),
1368 TapeFuncallArg::Str(_) => None,
1369 })
1370 .collect();
1371 for (k, &tk) in real_tape.iter().enumerate() {
1372 adj[tk] += w * derivs[k];
1373 let mut second_term = 0.0;
1374 for (l, &tl) in real_tape.iter().enumerate() {
1375 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1376 let h_kl = hes[lo + hi * (hi + 1) / 2];
1377 second_term += h_kl * dot[tl];
1378 }
1379 adj_dot[tk] += wd * derivs[k] + w * second_term;
1380 }
1381 }
1382 }
1383 }
1384 }
1385
1386 pub fn hessian_accumulate(
1393 &self,
1394 x: &[f64],
1395 weight: f64,
1396 hess_map: &HashMap<(usize, usize), usize>,
1397 values: &mut [f64],
1398 ) {
1399 let n = self.ops.len();
1400 if n == 0 || weight == 0.0 {
1401 return;
1402 }
1403 let v = self.forward(x);
1404 let var_indices = self.variables();
1405
1406 let mut dot = vec![0.0f64; n];
1413 let mut adj = vec![0.0f64; n];
1414 let mut adj_dot = vec![0.0f64; n];
1415 for &j in &var_indices {
1416 self.forward_tangent(&v, j, &mut dot);
1417
1418 adj.fill(0.0);
1421 adj_dot.fill(0.0);
1422 adj[n - 1] = 1.0;
1423
1424 for i in (0..n).rev() {
1425 let w = adj[i];
1426 let wd = adj_dot[i];
1427 if w == 0.0 && wd == 0.0 {
1428 continue;
1429 }
1430 match &self.ops[i] {
1431 TapeOp::Const(_) => {}
1432 TapeOp::Var(k) => {
1433 if wd != 0.0 && *k >= j {
1436 if let Some(&pos) = hess_map.get(&(*k, j)) {
1437 values[pos] += weight * wd;
1438 }
1439 }
1440 }
1441 TapeOp::Add(a, b) => {
1442 adj[*a] += w;
1443 adj[*b] += w;
1444 adj_dot[*a] += wd;
1445 adj_dot[*b] += wd;
1446 }
1447 TapeOp::Sub(a, b) => {
1448 adj[*a] += w;
1449 adj[*b] -= w;
1450 adj_dot[*a] += wd;
1451 adj_dot[*b] -= wd;
1452 }
1453 TapeOp::Mul(a, b) => {
1454 adj[*a] += w * v[*b];
1455 adj[*b] += w * v[*a];
1456 adj_dot[*a] += wd * v[*b] + w * dot[*b];
1457 adj_dot[*b] += wd * v[*a] + w * dot[*a];
1458 }
1459 TapeOp::Div(a, b) => {
1460 let vb = v[*b];
1463 let q = v[i];
1464 let qd = dot[i];
1465 adj[*a] += w / vb;
1466 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
1467 adj[*b] -= w * q / vb;
1468 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
1469 }
1470 TapeOp::Pow(a, b) => {
1471 let u = v[*a];
1472 let r = v[*b];
1473 let du = dot[*a];
1474 let dr = dot[*b];
1475 if r != 0.0 {
1476 if u != 0.0 {
1477 let p_a = r * u.powf(r - 1.0);
1478 adj[*a] += w * p_a;
1479 let mut dp_a = dr * u.powf(r - 1.0);
1480 if u > 0.0 {
1481 dp_a +=
1482 r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
1483 } else {
1484 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
1485 }
1486 adj_dot[*a] += wd * p_a + w * dp_a;
1487 } else if r >= 2.0 {
1488 let p_a = 0.0;
1489 adj[*a] += w * p_a;
1490 let dp_a = if r == 2.0 {
1491 2.0 * du
1492 } else {
1493 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
1494 };
1495 adj_dot[*a] += wd * p_a + w * dp_a;
1496 }
1497 }
1498 if u > 0.0 {
1499 let ln_u = u.ln();
1500 let p_b = v[i] * ln_u;
1501 adj[*b] += w * p_b;
1502 let dur = v[i] * (r * du / u + dr * ln_u);
1503 let dp_b = dur * ln_u + v[i] * du / u;
1504 adj_dot[*b] += wd * p_b + w * dp_b;
1505 }
1506 }
1507 TapeOp::Neg(a) => {
1508 adj[*a] -= w;
1509 adj_dot[*a] -= wd;
1510 }
1511 TapeOp::Abs(a) => {
1512 let s = if v[*a] >= 0.0 { 1.0 } else { -1.0 };
1513 adj[*a] += w * s;
1514 adj_dot[*a] += wd * s;
1515 }
1516 TapeOp::Sqrt(a) => {
1517 let sv = v[i];
1518 if sv > 0.0 {
1519 let fp = 0.5 / sv;
1520 let fpp = -0.25 / (v[*a] * sv);
1521 adj[*a] += w * fp;
1522 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
1523 }
1524 }
1525 TapeOp::Exp(a) => {
1526 let ev = v[i];
1527 adj[*a] += w * ev;
1528 adj_dot[*a] += wd * ev + w * ev * dot[*a];
1529 }
1530 TapeOp::Log(a) => {
1531 let u = v[*a];
1532 adj[*a] += w / u;
1533 adj_dot[*a] += wd / u + w * (-1.0 / (u * u)) * dot[*a];
1534 }
1535 TapeOp::Log10(a) => {
1536 let u = v[*a];
1537 let c = std::f64::consts::LN_10;
1538 adj[*a] += w / (u * c);
1539 adj_dot[*a] += wd / (u * c) + w * (-1.0 / (u * u * c)) * dot[*a];
1540 }
1541 TapeOp::Sin(a) => {
1542 let u = v[*a];
1543 let cu = u.cos();
1544 adj[*a] += w * cu;
1545 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
1546 }
1547 TapeOp::Cos(a) => {
1548 let u = v[*a];
1549 let su = u.sin();
1550 adj[*a] -= w * su;
1551 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
1552 }
1553 TapeOp::Tan(a) => {
1554 let t = v[i];
1555 let gp = 1.0 + t * t;
1556 let gpp = 2.0 * t * gp;
1557 adj[*a] += w * gp;
1558 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1559 }
1560 TapeOp::Atan(a) => {
1561 let u = v[*a];
1562 let d = 1.0 + u * u;
1563 let gp = 1.0 / d;
1564 let gpp = -2.0 * u / (d * d);
1565 adj[*a] += w * gp;
1566 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1567 }
1568 TapeOp::Acos(a) => {
1569 let u = v[*a];
1570 let s = 1.0 - u * u;
1571 let r = s.sqrt();
1572 let gp = -1.0 / r;
1573 let gpp = -u / (s * r);
1574 adj[*a] += w * gp;
1575 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1576 }
1577 TapeOp::Sinh(a) => {
1578 let u = v[*a];
1579 let gp = u.cosh();
1580 let gpp = u.sinh();
1581 adj[*a] += w * gp;
1582 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1583 }
1584 TapeOp::Cosh(a) => {
1585 let u = v[*a];
1586 let gp = u.sinh();
1587 let gpp = u.cosh();
1588 adj[*a] += w * gp;
1589 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1590 }
1591 TapeOp::Tanh(a) => {
1592 let t = v[i];
1593 let gp = 1.0 - t * t;
1594 let gpp = -2.0 * t * gp;
1595 adj[*a] += w * gp;
1596 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1597 }
1598 TapeOp::Asin(a) => {
1599 let u = v[*a];
1600 let s = 1.0 - u * u;
1601 let r = s.sqrt();
1602 let gp = 1.0 / r;
1603 let gpp = u / (s * r);
1604 adj[*a] += w * gp;
1605 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1606 }
1607 TapeOp::Acosh(a) => {
1608 let u = v[*a];
1609 let s = u * u - 1.0;
1610 let r = s.sqrt();
1611 let gp = 1.0 / r;
1612 let gpp = -u / (s * r);
1613 adj[*a] += w * gp;
1614 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1615 }
1616 TapeOp::Asinh(a) => {
1617 let u = v[*a];
1618 let s = u * u + 1.0;
1619 let r = s.sqrt();
1620 let gp = 1.0 / r;
1621 let gpp = -u / (s * r);
1622 adj[*a] += w * gp;
1623 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1624 }
1625 TapeOp::Atanh(a) => {
1626 let u = v[*a];
1627 let d = 1.0 - u * u;
1628 let gp = 1.0 / d;
1629 let gpp = 2.0 * u / (d * d);
1630 adj[*a] += w * gp;
1631 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1632 }
1633 TapeOp::Erf(a) => {
1634 let u = v[*a];
1635 let gp = erf_d1(u);
1636 let gpp = erf_d2(u);
1637 adj[*a] += w * gp;
1638 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1639 }
1640 TapeOp::XLogX(a) => {
1641 let u = v[*a];
1644 let gp = xlogx_d1(u);
1645 let gpp = xlogx_d2(u);
1646 adj[*a] += w * gp;
1647 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
1648 }
1649 TapeOp::CEntropy(a, b) => {
1650 let ua = v[*a];
1651 let ub = v[*b];
1652 let fa = centropy_da(ua, ub);
1653 let fb = centropy_db(ua, ub);
1654 let faa = centropy_daa(ua);
1655 let fab = centropy_dab(ub);
1656 let fbb = centropy_dbb(ua, ub);
1657 adj[*a] += w * fa;
1658 adj[*b] += w * fb;
1659 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1660 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1661 }
1662 TapeOp::Atan2(a, b) => {
1663 let y = v[*a];
1664 let x = v[*b];
1665 let d = y * y + x * x;
1666 let d2 = d * d;
1667 let fa = x / d;
1668 let fb = -y / d;
1669 let faa = -2.0 * y * x / d2;
1670 let fab = (y * y - x * x) / d2;
1671 let fbb = 2.0 * y * x / d2;
1672 adj[*a] += w * fa;
1673 adj[*b] += w * fb;
1674 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
1675 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
1676 }
1677 TapeOp::Min(a, b) => {
1681 let br = if v[*a] <= v[*b] { *a } else { *b };
1682 adj[br] += w;
1683 adj_dot[br] += wd;
1684 }
1685 TapeOp::Max(a, b) => {
1686 let br = if v[*a] >= v[*b] { *a } else { *b };
1687 adj[br] += w;
1688 adj_dot[br] += wd;
1689 }
1690 TapeOp::Cmp(_, _, _)
1692 | TapeOp::And(_, _)
1693 | TapeOp::Or(_, _)
1694 | TapeOp::Not(_) => {}
1695 TapeOp::Select(c, t, e) => {
1698 let br = if v[*c] != 0.0 { *t } else { *e };
1699 adj[br] += w;
1700 adj_dot[br] += wd;
1701 }
1702 TapeOp::Funcall(fc) => {
1703 let FuncallData { lib, name, args } = fc.as_ref();
1704 let call_args = funcall_to_ext_args(args, &v);
1705 let res = ext_eval_or_nan(lib, name, &call_args, args.len(), true, true);
1706 let derivs = res.derivs.expect("want_derivs=true returns derivs");
1707 let hes = res.hessian.expect("want_hes=true returns hessian");
1708 let real_tape: Vec<usize> = args
1709 .iter()
1710 .filter_map(|a| match a {
1711 TapeFuncallArg::Tape(t) => Some(*t),
1712 TapeFuncallArg::Str(_) => None,
1713 })
1714 .collect();
1715 for (k, &tk) in real_tape.iter().enumerate() {
1716 adj[tk] += w * derivs[k];
1717 let mut second_term = 0.0;
1718 for (l, &tl) in real_tape.iter().enumerate() {
1719 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
1720 let h_kl = hes[lo + hi * (hi + 1) / 2];
1721 second_term += h_kl * dot[tl];
1722 }
1723 adj_dot[tk] += wd * derivs[k] + w * second_term;
1724 }
1725 }
1726 }
1727 }
1728 }
1729 }
1730
1731 pub fn hessian_sparsity(&self) -> BTreeSet<(usize, usize)> {
1736 let n = self.ops.len();
1737 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
1738 let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
1739
1740 let mut last_use: Vec<usize> = (0..n).collect();
1755 for (i, op) in self.ops.iter().enumerate() {
1756 for_each_input(op, |a| last_use[a] = i);
1757 }
1758 macro_rules! merge {
1762 ($a:expr, $b:expr, $i:expr) => {{
1763 let (a, b, i) = ($a, $b, $i);
1764 if a == b {
1765 if last_use[a] == i {
1766 std::mem::take(&mut var_sets[a])
1767 } else {
1768 var_sets[a].clone()
1769 }
1770 } else {
1771 let take_a = last_use[a] == i
1773 && (last_use[b] != i || var_sets[a].len() >= var_sets[b].len());
1774 if take_a {
1775 let mut s = std::mem::take(&mut var_sets[a]);
1776 s.extend(var_sets[b].iter().copied());
1777 s
1778 } else if last_use[b] == i {
1779 let mut s = std::mem::take(&mut var_sets[b]);
1780 s.extend(var_sets[a].iter().copied());
1781 s
1782 } else {
1783 var_sets[a].union(&var_sets[b]).copied().collect()
1784 }
1785 }
1786 }};
1787 }
1788 macro_rules! carry {
1790 ($a:expr, $i:expr) => {{
1791 let (a, i) = ($a, $i);
1792 if last_use[a] == i {
1793 std::mem::take(&mut var_sets[a])
1794 } else {
1795 var_sets[a].clone()
1796 }
1797 }};
1798 }
1799
1800 let emit_cross =
1801 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1802 for &v1 in s1 {
1803 for &v2 in s2 {
1804 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
1805 pairs.insert((r, c));
1806 }
1807 }
1808 };
1809 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
1810 let vars: Vec<usize> = s.iter().copied().collect();
1811 for (ai, &vi) in vars.iter().enumerate() {
1812 for &vj in &vars[..=ai] {
1813 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
1814 pairs.insert((r, c));
1815 }
1816 }
1817 };
1818
1819 for (i, op) in self.ops.iter().enumerate() {
1820 let vset = match op {
1821 TapeOp::Const(_) => BTreeSet::new(),
1822 TapeOp::Var(j) => {
1823 let mut s = BTreeSet::new();
1824 s.insert(*j);
1825 s
1826 }
1827 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => merge!(*a, *b, i),
1828 TapeOp::Neg(a) | TapeOp::Abs(a) => carry!(*a, i),
1829 TapeOp::Mul(a, b) => {
1830 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1831 merge!(*a, *b, i)
1832 }
1833 TapeOp::Div(a, b) => {
1834 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
1835 emit_self(&var_sets[*b], &mut pairs);
1836 merge!(*a, *b, i)
1837 }
1838 TapeOp::Pow(a, b) => {
1839 let combined = merge!(*a, *b, i);
1840 emit_self(&combined, &mut pairs);
1841 combined
1842 }
1843 TapeOp::Sqrt(a)
1844 | TapeOp::Exp(a)
1845 | TapeOp::Log(a)
1846 | TapeOp::Log10(a)
1847 | TapeOp::Sin(a)
1848 | TapeOp::Cos(a)
1849 | TapeOp::Tan(a)
1850 | TapeOp::Atan(a)
1851 | TapeOp::Acos(a)
1852 | TapeOp::Sinh(a)
1853 | TapeOp::Cosh(a)
1854 | TapeOp::Tanh(a)
1855 | TapeOp::Asin(a)
1856 | TapeOp::Acosh(a)
1857 | TapeOp::Asinh(a)
1858 | TapeOp::Erf(a)
1859 | TapeOp::XLogX(a)
1860 | TapeOp::Atanh(a) => {
1861 emit_self(&var_sets[*a], &mut pairs);
1862 carry!(*a, i)
1863 }
1864 TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
1869 let combined = merge!(*a, *b, i);
1870 emit_self(&combined, &mut pairs);
1871 combined
1872 }
1873 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
1879 BTreeSet::new()
1880 }
1881 TapeOp::Select(_c, t, e) => merge!(*t, *e, i),
1888 TapeOp::Min(a, b) | TapeOp::Max(a, b) => merge!(*a, *b, i),
1895 TapeOp::Funcall(fc) => {
1896 let args = &fc.args;
1897 let mut combined: BTreeSet<usize> = BTreeSet::new();
1898 for arg in args {
1899 if let TapeFuncallArg::Tape(t) = arg {
1900 for &vv in &var_sets[*t] {
1901 combined.insert(vv);
1902 }
1903 }
1904 }
1905 emit_self(&combined, &mut pairs);
1906 combined
1907 }
1908 };
1909 for_each_input(op, |a| {
1915 if last_use[a] == i {
1916 var_sets[a].clear();
1917 }
1918 });
1919 var_sets.push(vset);
1920 }
1921 pairs
1922 }
1923}
1924
1925fn for_each_input(op: &TapeOp, mut f: impl FnMut(usize)) {
1936 match op {
1937 TapeOp::Const(_) | TapeOp::Var(_) => {}
1938 TapeOp::Neg(a)
1939 | TapeOp::Abs(a)
1940 | TapeOp::Sqrt(a)
1941 | TapeOp::Exp(a)
1942 | TapeOp::Log(a)
1943 | TapeOp::Log10(a)
1944 | TapeOp::Sin(a)
1945 | TapeOp::Cos(a)
1946 | TapeOp::Tan(a)
1947 | TapeOp::Atan(a)
1948 | TapeOp::Acos(a)
1949 | TapeOp::Sinh(a)
1950 | TapeOp::Cosh(a)
1951 | TapeOp::Tanh(a)
1952 | TapeOp::Asin(a)
1953 | TapeOp::Acosh(a)
1954 | TapeOp::Asinh(a)
1955 | TapeOp::Atanh(a)
1956 | TapeOp::Erf(a)
1957 | TapeOp::XLogX(a)
1958 | TapeOp::Not(a) => f(*a),
1959 TapeOp::Add(a, b)
1960 | TapeOp::Sub(a, b)
1961 | TapeOp::Mul(a, b)
1962 | TapeOp::Div(a, b)
1963 | TapeOp::Pow(a, b)
1964 | TapeOp::CEntropy(a, b)
1965 | TapeOp::Atan2(a, b)
1966 | TapeOp::Min(a, b)
1967 | TapeOp::Max(a, b)
1968 | TapeOp::And(a, b)
1969 | TapeOp::Or(a, b)
1970 | TapeOp::Cmp(_, a, b) => {
1971 f(*a);
1972 f(*b);
1973 }
1974 TapeOp::Select(c, t, e) => {
1975 f(*c);
1976 f(*t);
1977 f(*e);
1978 }
1979 TapeOp::Funcall(fc) => {
1980 for arg in &fc.args {
1981 if let TapeFuncallArg::Tape(t) = arg {
1982 f(*t);
1983 }
1984 }
1985 }
1986 }
1987}
1988
1989fn build_recursive(
1990 expr: &Expr,
1991 ops: &mut Vec<TapeOp>,
1992 cache: &mut HashMap<*const Expr, usize>,
1993 resolver: &ExternalResolver,
1994) -> usize {
1995 match expr {
1996 Expr::Const(c) => {
1997 let idx = ops.len();
1998 ops.push(TapeOp::Const(*c));
1999 idx
2000 }
2001 Expr::Var(i) => {
2002 let idx = ops.len();
2003 ops.push(TapeOp::Var(*i));
2004 idx
2005 }
2006 Expr::Binary(op, a, b) => {
2007 if let BinOp::Pow = op {
2015 if let Some(c) = peek_const(b) {
2016 if let Some(idx) = try_emit_const_pow(a, c, ops, cache, resolver) {
2017 return idx;
2018 }
2019 }
2020 }
2021 let l = build_recursive(a, ops, cache, resolver);
2022 let r = build_recursive(b, ops, cache, resolver);
2023 let idx = ops.len();
2024 ops.push(match op {
2025 BinOp::Add => TapeOp::Add(l, r),
2026 BinOp::Sub => TapeOp::Sub(l, r),
2027 BinOp::Mul => TapeOp::Mul(l, r),
2028 BinOp::Div => TapeOp::Div(l, r),
2029 BinOp::Pow => TapeOp::Pow(l, r),
2030 BinOp::Atan2 => TapeOp::Atan2(l, r),
2031 BinOp::CEntropy => TapeOp::CEntropy(l, r),
2032 });
2033 idx
2034 }
2035 Expr::Unary(op, a) => {
2036 let v = build_recursive(a, ops, cache, resolver);
2037 let idx = ops.len();
2038 ops.push(match op {
2039 UnaryOp::Neg => TapeOp::Neg(v),
2040 UnaryOp::Sqrt => TapeOp::Sqrt(v),
2041 UnaryOp::Log => TapeOp::Log(v),
2042 UnaryOp::Log10 => TapeOp::Log10(v),
2043 UnaryOp::Exp => TapeOp::Exp(v),
2044 UnaryOp::Abs => TapeOp::Abs(v),
2045 UnaryOp::Sin => TapeOp::Sin(v),
2046 UnaryOp::Cos => TapeOp::Cos(v),
2047 UnaryOp::Tan => TapeOp::Tan(v),
2048 UnaryOp::Atan => TapeOp::Atan(v),
2049 UnaryOp::Acos => TapeOp::Acos(v),
2050 UnaryOp::Sinh => TapeOp::Sinh(v),
2051 UnaryOp::Cosh => TapeOp::Cosh(v),
2052 UnaryOp::Tanh => TapeOp::Tanh(v),
2053 UnaryOp::Asin => TapeOp::Asin(v),
2054 UnaryOp::Acosh => TapeOp::Acosh(v),
2055 UnaryOp::Asinh => TapeOp::Asinh(v),
2056 UnaryOp::Atanh => TapeOp::Atanh(v),
2057 UnaryOp::Erf => TapeOp::Erf(v),
2058 UnaryOp::XLogX => TapeOp::XLogX(v),
2059 });
2060 idx
2061 }
2062 Expr::Sum(args) => {
2063 if args.is_empty() {
2064 let idx = ops.len();
2065 ops.push(TapeOp::Const(0.0));
2066 return idx;
2067 }
2068 let mut acc = build_recursive(&args[0], ops, cache, resolver);
2069 for a in &args[1..] {
2070 let next = build_recursive(a, ops, cache, resolver);
2071 let idx = ops.len();
2072 ops.push(TapeOp::Add(acc, next));
2073 acc = idx;
2074 }
2075 acc
2076 }
2077 Expr::MinList(args) | Expr::MaxList(args) => {
2085 let is_min = matches!(expr, Expr::MinList(_));
2086 if args.is_empty() {
2087 let idx = ops.len();
2088 ops.push(TapeOp::Const(0.0));
2089 return idx;
2090 }
2091 let mut acc = build_recursive(&args[0], ops, cache, resolver);
2092 for a in &args[1..] {
2093 let next = build_recursive(a, ops, cache, resolver);
2094 let idx = ops.len();
2095 ops.push(if is_min {
2096 TapeOp::Min(acc, next)
2097 } else {
2098 TapeOp::Max(acc, next)
2099 });
2100 acc = idx;
2101 }
2102 acc
2103 }
2104 Expr::Cse(body) => {
2105 let key = Arc::as_ptr(body) as *const Expr;
2112 if let Some(&idx) = cache.get(&key) {
2113 idx
2114 } else {
2115 let idx = build_recursive(body, ops, cache, resolver);
2116 cache.insert(key, idx);
2117 idx
2118 }
2119 }
2120 Expr::Compare(op, a, b) => {
2121 let l = build_recursive(a, ops, cache, resolver);
2122 let r = build_recursive(b, ops, cache, resolver);
2123 let idx = ops.len();
2124 ops.push(TapeOp::Cmp(*op, l, r));
2125 idx
2126 }
2127 Expr::And(a, b) => {
2128 let l = build_recursive(a, ops, cache, resolver);
2129 let r = build_recursive(b, ops, cache, resolver);
2130 let idx = ops.len();
2131 ops.push(TapeOp::And(l, r));
2132 idx
2133 }
2134 Expr::Or(a, b) => {
2135 let l = build_recursive(a, ops, cache, resolver);
2136 let r = build_recursive(b, ops, cache, resolver);
2137 let idx = ops.len();
2138 ops.push(TapeOp::Or(l, r));
2139 idx
2140 }
2141 Expr::Not(a) => {
2142 let v = build_recursive(a, ops, cache, resolver);
2143 let idx = ops.len();
2144 ops.push(TapeOp::Not(v));
2145 idx
2146 }
2147 Expr::Cond { cond, then_, else_ } => {
2148 let c = build_recursive(cond, ops, cache, resolver);
2149 let t = build_recursive(then_, ops, cache, resolver);
2150 let e = build_recursive(else_, ops, cache, resolver);
2151 let idx = ops.len();
2152 ops.push(TapeOp::Select(c, t, e));
2153 idx
2154 }
2155 Expr::Funcall { id, args } => {
2156 let (lib, name) = resolver
2157 .funcs_by_id
2158 .get(id)
2159 .unwrap_or_else(|| panic!("unresolved AMPL funcall id {id}"));
2160 let tape_args: Vec<TapeFuncallArg> = args
2161 .iter()
2162 .map(|a| match a {
2163 FuncallArg::Real(e) => {
2164 TapeFuncallArg::Tape(build_recursive(e, ops, cache, resolver))
2165 }
2166 FuncallArg::Str(s) => TapeFuncallArg::Str(s.clone()),
2167 })
2168 .collect();
2169 let idx = ops.len();
2170 ops.push(TapeOp::Funcall(Box::new(FuncallData {
2171 lib: Arc::clone(lib),
2172 name: name.clone(),
2173 args: tape_args,
2174 })));
2175 idx
2176 }
2177 }
2178}
2179
2180fn peek_const(e: &Expr) -> Option<f64> {
2184 match e {
2185 Expr::Const(c) => Some(*c),
2186 Expr::Cse(body) => peek_const(body),
2187 _ => None,
2188 }
2189}
2190
2191fn try_emit_const_pow(
2199 base_expr: &Expr,
2200 c: f64,
2201 ops: &mut Vec<TapeOp>,
2202 cache: &mut HashMap<*const Expr, usize>,
2203 resolver: &ExternalResolver,
2204) -> Option<usize> {
2205 if c == 0.0 {
2206 let idx = ops.len();
2207 ops.push(TapeOp::Const(1.0));
2208 return Some(idx);
2209 }
2210 if c == 1.0 {
2211 return Some(build_recursive(base_expr, ops, cache, resolver));
2212 }
2213 if c == 0.5 {
2214 let b = build_recursive(base_expr, ops, cache, resolver);
2215 let idx = ops.len();
2216 ops.push(TapeOp::Sqrt(b));
2217 return Some(idx);
2218 }
2219 if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
2224 let n = c.abs() as u32;
2225 if n == 0 {
2226 let idx = ops.len();
2228 ops.push(TapeOp::Const(1.0));
2229 return Some(idx);
2230 }
2231 let b = build_recursive(base_expr, ops, cache, resolver);
2232 let pos = emit_int_pow(b, n, ops);
2233 if c < 0.0 {
2234 let one_idx = ops.len();
2237 ops.push(TapeOp::Const(1.0));
2238 let idx = ops.len();
2239 ops.push(TapeOp::Div(one_idx, pos));
2240 return Some(idx);
2241 }
2242 return Some(pos);
2243 }
2244 None
2245}
2246
2247fn emit_int_pow(base: usize, n: u32, ops: &mut Vec<TapeOp>) -> usize {
2251 debug_assert!(n >= 1);
2252 if n == 1 {
2253 return base;
2254 }
2255 let half = emit_int_pow(base, n / 2, ops);
2256 let squared = ops.len();
2257 ops.push(TapeOp::Mul(half, half));
2258 if n % 2 == 1 {
2259 let idx = ops.len();
2260 ops.push(TapeOp::Mul(squared, base));
2261 idx
2262 } else {
2263 squared
2264 }
2265}
2266
2267#[derive(Debug, Clone)]
2295pub enum SummandOp {
2296 Local(TapeOp),
2299 Shared(usize),
2303}
2304
2305#[derive(Debug, Clone)]
2306pub struct Summand {
2307 pub ops: Vec<SummandOp>,
2308 pub root_slot: usize,
2310 pub local_reach: Vec<usize>,
2312 pub prelude_reach: Vec<usize>,
2315 pub local_vars: Vec<usize>,
2317 pub prelude_vars: Vec<usize>,
2319 pub all_vars: Vec<usize>,
2326}
2327
2328#[derive(Debug, Clone)]
2329pub struct HybridTape {
2330 pub prelude: Vec<TapeOp>,
2335 pub summands: Vec<Summand>,
2336}
2337
2338impl HybridTape {
2339 pub fn build_multi(exprs: &[Expr]) -> Self {
2344 let mut cse_count: HashMap<*const Expr, usize> = HashMap::new();
2348 for e in exprs {
2349 let mut seen_in_root: HashSet<*const Expr> = HashSet::new();
2350 count_cse_appearances(e, &mut seen_in_root, &mut cse_count);
2351 }
2352
2353 let mut prelude: Vec<TapeOp> = Vec::new();
2358 let mut prelude_map: HashMap<*const Expr, usize> = HashMap::new();
2359 let mut summands: Vec<Summand> = Vec::with_capacity(exprs.len());
2360 for e in exprs {
2361 let mut local: Vec<SummandOp> = Vec::new();
2362 let mut local_cache: HashMap<*const Expr, usize> = HashMap::new();
2363 let root_slot = build_into_summand(
2364 e,
2365 &mut local,
2366 &mut local_cache,
2367 &mut prelude,
2368 &mut prelude_map,
2369 &cse_count,
2370 );
2371 summands.push(Summand {
2372 ops: local,
2373 root_slot,
2374 local_reach: Vec::new(),
2375 prelude_reach: Vec::new(),
2376 local_vars: Vec::new(),
2377 prelude_vars: Vec::new(),
2378 all_vars: Vec::new(),
2379 });
2380 }
2381
2382 let mut p_visited: Vec<u32> = vec![0; prelude.len()];
2386 let mut p_epoch: u32 = 0;
2387 let mut p_stack: Vec<usize> = Vec::new();
2388 for s in &mut summands {
2389 let (local_reach, shared_refs) = compute_local_reach(&s.ops, s.root_slot);
2390 s.local_reach = local_reach;
2391
2392 let mut lv: BTreeSet<usize> = BTreeSet::new();
2393 for &i in &s.local_reach {
2394 if let SummandOp::Local(TapeOp::Var(j)) = &s.ops[i] {
2395 lv.insert(*j);
2396 }
2397 }
2398 s.local_vars = lv.iter().copied().collect();
2399
2400 if !shared_refs.is_empty() {
2401 p_epoch += 1;
2402 let mut preach: Vec<usize> = Vec::new();
2403 for &start in &shared_refs {
2404 bfs_prelude(
2405 &prelude,
2406 start,
2407 &mut p_visited,
2408 p_epoch,
2409 &mut p_stack,
2410 &mut preach,
2411 );
2412 }
2413 preach.sort_unstable();
2414 s.prelude_vars = vars_in(&prelude, &preach);
2415 s.prelude_reach = preach;
2416 }
2417
2418 let mut av: BTreeSet<usize> = lv;
2419 for &v in &s.prelude_vars {
2420 av.insert(v);
2421 }
2422 s.all_vars = av.into_iter().collect();
2423 }
2424
2425 HybridTape { prelude, summands }
2426 }
2427
2428 pub fn n_prelude_ops(&self) -> usize {
2429 self.prelude.len()
2430 }
2431 pub fn n_summands(&self) -> usize {
2432 self.summands.len()
2433 }
2434 pub fn max_summand_ops(&self) -> usize {
2435 self.summands.iter().map(|s| s.ops.len()).max().unwrap_or(0)
2436 }
2437 pub fn total_local_ops(&self) -> usize {
2438 self.summands.iter().map(|s| s.ops.len()).sum()
2439 }
2440
2441 pub fn forward_prelude(&self, x: &[f64], prelude_vals: &mut [f64]) {
2444 debug_assert_eq!(prelude_vals.len(), self.prelude.len());
2445 for i in 0..self.prelude.len() {
2446 prelude_vals[i] = fwd_step(&self.prelude[i], x, prelude_vals);
2447 }
2448 }
2449
2450 pub fn forward_summand(
2453 &self,
2454 s: &Summand,
2455 x: &[f64],
2456 prelude_vals: &[f64],
2457 local_vals: &mut [f64],
2458 ) {
2459 debug_assert!(local_vals.len() >= s.ops.len());
2460 for i in 0..s.ops.len() {
2461 local_vals[i] = match &s.ops[i] {
2462 SummandOp::Local(op) => fwd_step(op, x, local_vals),
2463 SummandOp::Shared(k) => prelude_vals[*k],
2464 };
2465 }
2466 }
2467
2468 #[inline]
2470 pub fn root_value(&self, s: &Summand, local_vals: &[f64]) -> f64 {
2471 local_vals[s.root_slot]
2472 }
2473
2474 #[allow(clippy::too_many_arguments)]
2481 pub fn gradient_summand(
2482 &self,
2483 s: &Summand,
2484 prelude_vals: &[f64],
2485 local_vals: &[f64],
2486 seed: f64,
2487 grad: &mut [f64],
2488 local_adj: &mut [f64],
2489 prelude_adj: &mut [f64],
2490 ) {
2491 if seed == 0.0 || s.local_reach.is_empty() {
2492 return;
2493 }
2494 for &i in &s.local_reach {
2495 local_adj[i] = 0.0;
2496 }
2497 for &i in &s.prelude_reach {
2498 prelude_adj[i] = 0.0;
2499 }
2500 local_adj[s.root_slot] = seed;
2501 for &i in s.local_reach.iter().rev() {
2502 let a = local_adj[i];
2503 if a == 0.0 {
2504 continue;
2505 }
2506 match &s.ops[i] {
2507 SummandOp::Local(op) => rev_step(op, i, local_vals, local_adj, a, grad),
2508 SummandOp::Shared(k) => {
2509 prelude_adj[*k] += a;
2510 }
2511 }
2512 }
2513 for &i in s.prelude_reach.iter().rev() {
2514 let a = prelude_adj[i];
2515 if a == 0.0 {
2516 continue;
2517 }
2518 rev_step(&self.prelude[i], i, prelude_vals, prelude_adj, a, grad);
2519 }
2520 }
2521
2522 pub fn prelude_tangent(
2542 &self,
2543 prelude_vals: &[f64],
2544 seed: &[f64],
2545 reach: &[u32],
2546 prelude_dot: &mut [f64],
2547 ) {
2548 debug_assert!(prelude_dot.len() >= self.prelude.len());
2549 for &i in reach {
2550 let i = i as usize;
2551 prelude_dot[i] = fwd_dir_step(&self.prelude[i], seed, prelude_vals, prelude_dot, i);
2552 }
2553 }
2554
2555 #[allow(clippy::too_many_arguments)]
2583 pub fn hessian_summand_directional(
2584 &self,
2585 s: &Summand,
2586 local_vals: &[f64],
2587 prelude_dot: &[f64],
2588 seed: &[f64],
2589 weight: f64,
2590 out: &mut [f64],
2591 local_dot: &mut [f64],
2592 local_adj: &mut [f64],
2593 local_adj_dot: &mut [f64],
2594 prelude_adj: &mut [f64],
2595 prelude_adj_dot: &mut [f64],
2596 ) {
2597 if weight == 0.0 || s.local_reach.is_empty() {
2598 return;
2599 }
2600 for &i in &s.local_reach {
2601 local_adj[i] = 0.0;
2602 local_adj_dot[i] = 0.0;
2603 }
2604 for &i in &s.local_reach {
2608 local_dot[i] = match &s.ops[i] {
2609 SummandOp::Local(op) => fwd_dir_step(op, seed, local_vals, local_dot, i),
2610 SummandOp::Shared(k) => prelude_dot[*k],
2611 };
2612 }
2613 local_adj[s.root_slot] = 1.0;
2614 for &i in s.local_reach.iter().rev() {
2615 let w = local_adj[i];
2616 let wd = local_adj_dot[i];
2617 if w == 0.0 && wd == 0.0 {
2618 continue;
2619 }
2620 match &s.ops[i] {
2621 SummandOp::Local(op) => {
2622 ror_dir_step(
2623 op,
2624 i,
2625 local_vals,
2626 local_dot,
2627 local_adj,
2628 local_adj_dot,
2629 w,
2630 wd,
2631 weight,
2632 out,
2633 );
2634 }
2635 SummandOp::Shared(k) => {
2636 prelude_adj[*k] += weight * w;
2637 prelude_adj_dot[*k] += weight * wd;
2638 }
2639 }
2640 }
2641 }
2642
2643 pub fn prelude_reverse_directional(
2657 &self,
2658 prelude_vals: &[f64],
2659 prelude_dot: &[f64],
2660 reach: &[u32],
2661 out: &mut [f64],
2662 prelude_adj: &mut [f64],
2663 prelude_adj_dot: &mut [f64],
2664 ) {
2665 for &i in reach.iter().rev() {
2666 let i = i as usize;
2667 let w = prelude_adj[i];
2668 let wd = prelude_adj_dot[i];
2669 if w == 0.0 && wd == 0.0 {
2670 continue;
2671 }
2672 prelude_adj[i] = 0.0;
2673 prelude_adj_dot[i] = 0.0;
2674 ror_dir_step(
2675 &self.prelude[i],
2676 i,
2677 prelude_vals,
2678 prelude_dot,
2679 prelude_adj,
2680 prelude_adj_dot,
2681 w,
2682 wd,
2683 1.0,
2684 out,
2685 );
2686 }
2687 }
2688
2689 pub fn hessian_sparsity_all(&self) -> BTreeSet<(usize, usize)> {
2692 let mut pairs = hessian_sparsity_impl(&self.prelude);
2693
2694 let prelude_var_sets = compute_var_sets(&self.prelude);
2697
2698 for s in &self.summands {
2699 summand_sparsity(&s.ops, &prelude_var_sets, &mut pairs);
2700 }
2701 pairs
2702 }
2703}
2704
2705pub fn hybrid_supported(exprs: &[Expr]) -> bool {
2731 let mut stack: Vec<&Expr> = exprs.iter().collect();
2732 let mut seen_cse: HashSet<*const Expr> = HashSet::new();
2733 while let Some(e) = stack.pop() {
2734 match e {
2735 Expr::Const(_) | Expr::Var(_) => {}
2736 Expr::Binary(_, a, b) => {
2737 stack.push(a);
2738 stack.push(b);
2739 }
2740 Expr::Unary(_, a) => stack.push(a),
2741 Expr::Sum(args) => stack.extend(args.iter()),
2742 Expr::Cse(body) => {
2743 if seen_cse.insert(Arc::as_ptr(body)) {
2744 stack.push(body);
2745 }
2746 }
2747 Expr::Compare(..)
2748 | Expr::And(..)
2749 | Expr::Or(..)
2750 | Expr::Not(_)
2751 | Expr::Cond { .. }
2752 | Expr::MinList(_)
2753 | Expr::MaxList(_)
2754 | Expr::Funcall { .. } => return false,
2755 }
2756 }
2757 true
2758}
2759
2760fn cse_contains_funcall(expr: &Expr) -> bool {
2761 match expr {
2762 Expr::Funcall { .. } => true,
2763 Expr::Const(_) | Expr::Var(_) => false,
2764 Expr::Binary(_, a, b) => cse_contains_funcall(a) || cse_contains_funcall(b),
2765 Expr::Unary(_, a) => cse_contains_funcall(a),
2766 Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2767 args.iter().any(cse_contains_funcall)
2768 }
2769 Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2770 cse_contains_funcall(a) || cse_contains_funcall(b)
2771 }
2772 Expr::Not(a) => cse_contains_funcall(a),
2773 Expr::Cond { cond, then_, else_ } => {
2774 cse_contains_funcall(cond) || cse_contains_funcall(then_) || cse_contains_funcall(else_)
2775 }
2776 Expr::Cse(body) => cse_contains_funcall(body),
2777 }
2778}
2779
2780fn count_cse_appearances(
2781 e: &Expr,
2782 seen_in_root: &mut HashSet<*const Expr>,
2783 counts: &mut HashMap<*const Expr, usize>,
2784) {
2785 match e {
2786 Expr::Const(_) | Expr::Var(_) => {}
2787 Expr::Binary(_, a, b) => {
2788 count_cse_appearances(a, seen_in_root, counts);
2789 count_cse_appearances(b, seen_in_root, counts);
2790 }
2791 Expr::Unary(_, a) => count_cse_appearances(a, seen_in_root, counts),
2792 Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
2793 for a in args {
2794 count_cse_appearances(a, seen_in_root, counts);
2795 }
2796 }
2797 Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
2798 count_cse_appearances(a, seen_in_root, counts);
2799 count_cse_appearances(b, seen_in_root, counts);
2800 }
2801 Expr::Not(a) => count_cse_appearances(a, seen_in_root, counts),
2802 Expr::Cond { cond, then_, else_ } => {
2803 count_cse_appearances(cond, seen_in_root, counts);
2804 count_cse_appearances(then_, seen_in_root, counts);
2805 count_cse_appearances(else_, seen_in_root, counts);
2806 }
2807 Expr::Cse(body) => {
2808 let key = Arc::as_ptr(body) as *const Expr;
2809 if seen_in_root.insert(key) {
2810 *counts.entry(key).or_insert(0) += 1;
2811 count_cse_appearances(body, seen_in_root, counts);
2812 }
2813 }
2814 Expr::Funcall { args, .. } => {
2815 for arg in args {
2816 if let FuncallArg::Real(e) = arg {
2817 count_cse_appearances(e, seen_in_root, counts);
2818 }
2819 }
2820 }
2821 }
2822}
2823
2824fn build_into_summand(
2830 expr: &Expr,
2831 local: &mut Vec<SummandOp>,
2832 local_cache: &mut HashMap<*const Expr, usize>,
2833 prelude: &mut Vec<TapeOp>,
2834 prelude_map: &mut HashMap<*const Expr, usize>,
2835 cse_count: &HashMap<*const Expr, usize>,
2836) -> usize {
2837 match expr {
2838 Expr::Const(c) => {
2839 let i = local.len();
2840 local.push(SummandOp::Local(TapeOp::Const(*c)));
2841 i
2842 }
2843 Expr::Var(j) => {
2844 let i = local.len();
2845 local.push(SummandOp::Local(TapeOp::Var(*j)));
2846 i
2847 }
2848 Expr::Binary(op, a, b) => {
2849 if let BinOp::Pow = op {
2850 if let Some(c) = peek_const(b) {
2851 if let Some(i) = try_emit_const_pow_summand(
2852 a,
2853 c,
2854 local,
2855 local_cache,
2856 prelude,
2857 prelude_map,
2858 cse_count,
2859 ) {
2860 return i;
2861 }
2862 }
2863 }
2864 let l = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
2865 let r = build_into_summand(b, local, local_cache, prelude, prelude_map, cse_count);
2866 let i = local.len();
2867 local.push(SummandOp::Local(match op {
2868 BinOp::Add => TapeOp::Add(l, r),
2869 BinOp::Sub => TapeOp::Sub(l, r),
2870 BinOp::Mul => TapeOp::Mul(l, r),
2871 BinOp::Div => TapeOp::Div(l, r),
2872 BinOp::Pow => TapeOp::Pow(l, r),
2873 BinOp::Atan2 => TapeOp::Atan2(l, r),
2874 BinOp::CEntropy => TapeOp::CEntropy(l, r),
2875 }));
2876 i
2877 }
2878 Expr::Unary(op, a) => {
2879 let v = build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
2880 let i = local.len();
2881 local.push(SummandOp::Local(match op {
2882 UnaryOp::Neg => TapeOp::Neg(v),
2883 UnaryOp::Sqrt => TapeOp::Sqrt(v),
2884 UnaryOp::Log => TapeOp::Log(v),
2885 UnaryOp::Log10 => TapeOp::Log10(v),
2886 UnaryOp::Exp => TapeOp::Exp(v),
2887 UnaryOp::Abs => TapeOp::Abs(v),
2888 UnaryOp::Sin => TapeOp::Sin(v),
2889 UnaryOp::Cos => TapeOp::Cos(v),
2890 UnaryOp::Tan => TapeOp::Tan(v),
2891 UnaryOp::Atan => TapeOp::Atan(v),
2892 UnaryOp::Acos => TapeOp::Acos(v),
2893 UnaryOp::Sinh => TapeOp::Sinh(v),
2894 UnaryOp::Cosh => TapeOp::Cosh(v),
2895 UnaryOp::Tanh => TapeOp::Tanh(v),
2896 UnaryOp::Asin => TapeOp::Asin(v),
2897 UnaryOp::Acosh => TapeOp::Acosh(v),
2898 UnaryOp::Asinh => TapeOp::Asinh(v),
2899 UnaryOp::Atanh => TapeOp::Atanh(v),
2900 UnaryOp::Erf => TapeOp::Erf(v),
2901 UnaryOp::XLogX => TapeOp::XLogX(v),
2902 }));
2903 i
2904 }
2905 Expr::Sum(args) => {
2906 if args.is_empty() {
2907 let i = local.len();
2908 local.push(SummandOp::Local(TapeOp::Const(0.0)));
2909 return i;
2910 }
2911 let mut acc = build_into_summand(
2912 &args[0],
2913 local,
2914 local_cache,
2915 prelude,
2916 prelude_map,
2917 cse_count,
2918 );
2919 for a in &args[1..] {
2920 let nxt =
2921 build_into_summand(a, local, local_cache, prelude, prelude_map, cse_count);
2922 let i = local.len();
2923 local.push(SummandOp::Local(TapeOp::Add(acc, nxt)));
2924 acc = i;
2925 }
2926 acc
2927 }
2928 Expr::Cse(body) => {
2929 let key = Arc::as_ptr(body) as *const Expr;
2930 if let Some(&li) = local_cache.get(&key) {
2931 return li;
2932 }
2933 let promoted = cse_count.get(&key).copied().unwrap_or(0) >= 2;
2934 if promoted {
2935 if cse_contains_funcall(body) {
2942 panic!(
2943 "HybridTape: AMPL external function calls are not supported on the \
2944 hybrid (partial-separability) tape path. Build with \
2945 Tape::build_with_externals instead."
2946 );
2947 }
2948 let pslot =
2953 build_recursive(expr, prelude, prelude_map, &ExternalResolver::default());
2954 let li = local.len();
2955 local.push(SummandOp::Shared(pslot));
2956 local_cache.insert(key, li);
2957 li
2958 } else {
2959 let li =
2960 build_into_summand(body, local, local_cache, prelude, prelude_map, cse_count);
2961 local_cache.insert(key, li);
2962 li
2963 }
2964 }
2965 Expr::Compare(_, _, _)
2966 | Expr::And(_, _)
2967 | Expr::Or(_, _)
2968 | Expr::Not(_)
2969 | Expr::Cond { .. }
2970 | Expr::MinList(_)
2971 | Expr::MaxList(_) => {
2972 panic!(
2973 "HybridTape: conditional / logical / min-max opcodes (comparisons, \
2974 AND/OR/NOT, if-then-else, min/max lists) are not supported on the \
2975 hybrid (partial-separability) tape path. Build with \
2976 Tape::build_with_externals instead."
2977 );
2978 }
2979 Expr::Funcall { .. } => {
2980 panic!(
2981 "HybridTape: AMPL external function calls are not supported on the \
2982 hybrid (partial-separability) tape path. Build with Tape::build_with_externals \
2983 instead."
2984 );
2985 }
2986 }
2987}
2988
2989fn try_emit_const_pow_summand(
2992 base_expr: &Expr,
2993 c: f64,
2994 local: &mut Vec<SummandOp>,
2995 local_cache: &mut HashMap<*const Expr, usize>,
2996 prelude: &mut Vec<TapeOp>,
2997 prelude_map: &mut HashMap<*const Expr, usize>,
2998 cse_count: &HashMap<*const Expr, usize>,
2999) -> Option<usize> {
3000 if c == 0.0 {
3001 let i = local.len();
3002 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3003 return Some(i);
3004 }
3005 if c == 1.0 {
3006 return Some(build_into_summand(
3007 base_expr,
3008 local,
3009 local_cache,
3010 prelude,
3011 prelude_map,
3012 cse_count,
3013 ));
3014 }
3015 if c == 0.5 {
3016 let b = build_into_summand(
3017 base_expr,
3018 local,
3019 local_cache,
3020 prelude,
3021 prelude_map,
3022 cse_count,
3023 );
3024 let i = local.len();
3025 local.push(SummandOp::Local(TapeOp::Sqrt(b)));
3026 return Some(i);
3027 }
3028 if c.is_finite() && c.fract() == 0.0 && c.abs() <= 8.0 {
3029 let n = c.abs() as u32;
3030 if n == 0 {
3031 let i = local.len();
3032 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3033 return Some(i);
3034 }
3035 let b = build_into_summand(
3036 base_expr,
3037 local,
3038 local_cache,
3039 prelude,
3040 prelude_map,
3041 cse_count,
3042 );
3043 let pos = emit_int_pow_summand(b, n, local);
3044 if c < 0.0 {
3045 let one_idx = local.len();
3046 local.push(SummandOp::Local(TapeOp::Const(1.0)));
3047 let i = local.len();
3048 local.push(SummandOp::Local(TapeOp::Div(one_idx, pos)));
3049 return Some(i);
3050 }
3051 return Some(pos);
3052 }
3053 None
3054}
3055
3056fn emit_int_pow_summand(base: usize, n: u32, local: &mut Vec<SummandOp>) -> usize {
3057 debug_assert!(n >= 1);
3058 if n == 1 {
3059 return base;
3060 }
3061 let half = emit_int_pow_summand(base, n / 2, local);
3062 let squared = local.len();
3063 local.push(SummandOp::Local(TapeOp::Mul(half, half)));
3064 if n % 2 == 1 {
3065 let i = local.len();
3066 local.push(SummandOp::Local(TapeOp::Mul(squared, base)));
3067 i
3068 } else {
3069 squared
3070 }
3071}
3072
3073fn compute_local_reach(ops: &[SummandOp], root: usize) -> (Vec<usize>, Vec<usize>) {
3077 let mut visited = vec![false; ops.len()];
3078 let mut reach: Vec<usize> = Vec::new();
3079 let mut shared: BTreeSet<usize> = BTreeSet::new();
3080 let mut stack: Vec<usize> = Vec::with_capacity(16);
3081 visited[root] = true;
3082 reach.push(root);
3083 stack.push(root);
3084 while let Some(s) = stack.pop() {
3085 match &ops[s] {
3086 SummandOp::Local(op) => {
3087 let (a, b) = op_operands(op);
3088 if let Some(a) = a {
3089 if !visited[a] {
3090 visited[a] = true;
3091 reach.push(a);
3092 stack.push(a);
3093 }
3094 }
3095 if let Some(b) = b {
3096 if !visited[b] {
3097 visited[b] = true;
3098 reach.push(b);
3099 stack.push(b);
3100 }
3101 }
3102 }
3103 SummandOp::Shared(k) => {
3104 shared.insert(*k);
3105 }
3106 }
3107 }
3108 reach.sort_unstable();
3109 (reach, shared.into_iter().collect())
3110}
3111
3112fn bfs_prelude(
3116 prelude: &[TapeOp],
3117 start: usize,
3118 visited: &mut [u32],
3119 cur: u32,
3120 stack: &mut Vec<usize>,
3121 out: &mut Vec<usize>,
3122) {
3123 if visited[start] == cur {
3124 return;
3125 }
3126 visited[start] = cur;
3127 out.push(start);
3128 stack.push(start);
3129 while let Some(s) = stack.pop() {
3130 let (a, b) = op_operands(&prelude[s]);
3131 if let Some(a) = a {
3132 if visited[a] != cur {
3133 visited[a] = cur;
3134 out.push(a);
3135 stack.push(a);
3136 }
3137 }
3138 if let Some(b) = b {
3139 if visited[b] != cur {
3140 visited[b] = cur;
3141 out.push(b);
3142 stack.push(b);
3143 }
3144 }
3145 }
3146}
3147
3148fn compute_var_sets(ops: &[TapeOp]) -> Vec<BTreeSet<usize>> {
3152 let mut out: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3153 for op in ops {
3154 let vs: BTreeSet<usize> = match op {
3155 TapeOp::Const(_) => BTreeSet::new(),
3156 TapeOp::Var(j) => {
3157 let mut s = BTreeSet::new();
3158 s.insert(*j);
3159 s
3160 }
3161 TapeOp::Add(a, b)
3162 | TapeOp::Sub(a, b)
3163 | TapeOp::Mul(a, b)
3164 | TapeOp::Div(a, b)
3165 | TapeOp::Pow(a, b)
3166 | TapeOp::Atan2(a, b)
3167 | TapeOp::CEntropy(a, b) => out[*a].union(&out[*b]).copied().collect(),
3168 TapeOp::Neg(a)
3169 | TapeOp::Abs(a)
3170 | TapeOp::Sqrt(a)
3171 | TapeOp::Exp(a)
3172 | TapeOp::Log(a)
3173 | TapeOp::Log10(a)
3174 | TapeOp::Sin(a)
3175 | TapeOp::Cos(a)
3176 | TapeOp::Tan(a)
3177 | TapeOp::Atan(a)
3178 | TapeOp::Acos(a)
3179 | TapeOp::Sinh(a)
3180 | TapeOp::Cosh(a)
3181 | TapeOp::Tanh(a)
3182 | TapeOp::Asin(a)
3183 | TapeOp::Acosh(a)
3184 | TapeOp::Asinh(a)
3185 | TapeOp::Erf(a)
3186 | TapeOp::XLogX(a)
3187 | TapeOp::Atanh(a) => out[*a].clone(),
3188 TapeOp::Cmp(_, _, _)
3189 | TapeOp::And(_, _)
3190 | TapeOp::Or(_, _)
3191 | TapeOp::Not(_)
3192 | TapeOp::Select(_, _, _)
3193 | TapeOp::Min(_, _)
3194 | TapeOp::Max(_, _) => unreachable!(
3195 "HybridTape prelude cannot contain conditional / logical / min-max \
3196 TapeOps; build_into_summand panics on those Expr variants."
3197 ),
3198 TapeOp::Funcall(_) => unreachable!(
3199 "HybridTape prelude cannot contain TapeOp::Funcall; \
3200 build_into_summand panics on Expr::Funcall."
3201 ),
3202 };
3203 out.push(vs);
3204 }
3205 out
3206}
3207
3208fn summand_sparsity(
3213 ops: &[SummandOp],
3214 prelude_var_sets: &[BTreeSet<usize>],
3215 pairs: &mut BTreeSet<(usize, usize)>,
3216) {
3217 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(ops.len());
3218 let emit_cross =
3219 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3220 for &v1 in s1 {
3221 for &v2 in s2 {
3222 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
3223 pairs.insert((r, c));
3224 }
3225 }
3226 };
3227 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
3228 let vars: Vec<usize> = s.iter().copied().collect();
3229 for (ai, &vi) in vars.iter().enumerate() {
3230 for &vj in &vars[..=ai] {
3231 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
3232 pairs.insert((r, c));
3233 }
3234 }
3235 };
3236 for so in ops {
3237 let vset: BTreeSet<usize> = match so {
3238 SummandOp::Shared(k) => prelude_var_sets[*k].clone(),
3239 SummandOp::Local(op) => match op {
3240 TapeOp::Const(_) => BTreeSet::new(),
3241 TapeOp::Var(j) => {
3242 let mut s = BTreeSet::new();
3243 s.insert(*j);
3244 s
3245 }
3246 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
3247 var_sets[*a].union(&var_sets[*b]).copied().collect()
3248 }
3249 TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
3250 TapeOp::Mul(a, b) => {
3251 emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3252 var_sets[*a].union(&var_sets[*b]).copied().collect()
3253 }
3254 TapeOp::Div(a, b) => {
3255 emit_cross(&var_sets[*a], &var_sets[*b], pairs);
3256 emit_self(&var_sets[*b], pairs);
3257 var_sets[*a].union(&var_sets[*b]).copied().collect()
3258 }
3259 TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
3260 let combined: BTreeSet<usize> =
3261 var_sets[*a].union(&var_sets[*b]).copied().collect();
3262 emit_self(&combined, pairs);
3263 combined
3264 }
3265 TapeOp::Sqrt(a)
3266 | TapeOp::Exp(a)
3267 | TapeOp::Log(a)
3268 | TapeOp::Log10(a)
3269 | TapeOp::Sin(a)
3270 | TapeOp::Cos(a)
3271 | TapeOp::Tan(a)
3272 | TapeOp::Atan(a)
3273 | TapeOp::Acos(a)
3274 | TapeOp::Sinh(a)
3275 | TapeOp::Cosh(a)
3276 | TapeOp::Tanh(a)
3277 | TapeOp::Asin(a)
3278 | TapeOp::Acosh(a)
3279 | TapeOp::Asinh(a)
3280 | TapeOp::Erf(a)
3281 | TapeOp::XLogX(a)
3282 | TapeOp::Atanh(a) => {
3283 emit_self(&var_sets[*a], pairs);
3284 var_sets[*a].clone()
3285 }
3286 TapeOp::Cmp(_, _, _)
3287 | TapeOp::And(_, _)
3288 | TapeOp::Or(_, _)
3289 | TapeOp::Not(_)
3290 | TapeOp::Select(_, _, _)
3291 | TapeOp::Min(_, _)
3292 | TapeOp::Max(_, _) => unreachable!(
3293 "HybridTape summand cannot contain conditional / logical / min-max \
3294 TapeOps; build_into_summand panics on those Expr variants."
3295 ),
3296 TapeOp::Funcall(_) => unreachable!(
3297 "HybridTape summand cannot contain TapeOp::Funcall; \
3298 build_into_summand panics on Expr::Funcall."
3299 ),
3300 },
3301 };
3302 var_sets.push(vset);
3303 }
3304}
3305
3306#[inline]
3309pub(crate) fn op_operands(op: &TapeOp) -> (Option<usize>, Option<usize>) {
3310 match op {
3311 TapeOp::Const(_) | TapeOp::Var(_) => (None, None),
3312 TapeOp::Add(a, b)
3313 | TapeOp::Sub(a, b)
3314 | TapeOp::Mul(a, b)
3315 | TapeOp::Div(a, b)
3316 | TapeOp::Pow(a, b)
3317 | TapeOp::Atan2(a, b)
3318 | TapeOp::CEntropy(a, b) => (Some(*a), Some(*b)),
3319 TapeOp::Neg(a)
3320 | TapeOp::Abs(a)
3321 | TapeOp::Sqrt(a)
3322 | TapeOp::Exp(a)
3323 | TapeOp::Log(a)
3324 | TapeOp::Log10(a)
3325 | TapeOp::Sin(a)
3326 | TapeOp::Cos(a)
3327 | TapeOp::Tan(a)
3328 | TapeOp::Atan(a)
3329 | TapeOp::Acos(a)
3330 | TapeOp::Sinh(a)
3331 | TapeOp::Cosh(a)
3332 | TapeOp::Tanh(a)
3333 | TapeOp::Asin(a)
3334 | TapeOp::Acosh(a)
3335 | TapeOp::Asinh(a)
3336 | TapeOp::Erf(a)
3337 | TapeOp::XLogX(a)
3338 | TapeOp::Atanh(a) => (Some(*a), None),
3339 TapeOp::Cmp(_, a, b) | TapeOp::And(a, b) | TapeOp::Or(a, b) => (Some(*a), Some(*b)),
3345 TapeOp::Not(a) => (Some(*a), None),
3346 TapeOp::Select(_, _, _) => unreachable!(
3347 "op_operands: TapeOp::Select has three operands and is unsupported on \
3348 the HybridTape path"
3349 ),
3350 TapeOp::Min(_, _) | TapeOp::Max(_, _) => unreachable!(
3351 "op_operands: TapeOp::Min/Max are unsupported on the HybridTape path \
3352 (build_into_summand rejects min/max lists)"
3353 ),
3354 TapeOp::Funcall(_) => unreachable!(
3362 "op_operands: TapeOp::Funcall is unsupported on the HybridTape path \
3363 (build_into_summand rejects external function calls)"
3364 ),
3365 }
3366}
3367
3368fn vars_in(ops: &[TapeOp], reach: &[usize]) -> Vec<usize> {
3369 let mut s: BTreeSet<usize> = BTreeSet::new();
3370 for &i in reach {
3371 if let TapeOp::Var(j) = &ops[i] {
3372 s.insert(*j);
3373 }
3374 }
3375 s.into_iter().collect()
3376}
3377
3378#[inline]
3381fn fwd_step(op: &TapeOp, x: &[f64], vals: &[f64]) -> f64 {
3382 match op {
3383 TapeOp::Const(c) => *c,
3384 TapeOp::Var(i) => x[*i],
3385 TapeOp::Add(a, b) => vals[*a] + vals[*b],
3386 TapeOp::Sub(a, b) => vals[*a] - vals[*b],
3387 TapeOp::Mul(a, b) => vals[*a] * vals[*b],
3388 TapeOp::Div(a, b) => vals[*a] / vals[*b],
3389 TapeOp::Pow(a, b) => vals[*a].powf(vals[*b]),
3390 TapeOp::Neg(a) => -vals[*a],
3391 TapeOp::Abs(a) => vals[*a].abs(),
3392 TapeOp::Sqrt(a) => vals[*a].sqrt(),
3393 TapeOp::Exp(a) => vals[*a].exp(),
3394 TapeOp::Log(a) => vals[*a].ln(),
3395 TapeOp::Log10(a) => vals[*a].log10(),
3396 TapeOp::Sin(a) => vals[*a].sin(),
3397 TapeOp::Cos(a) => vals[*a].cos(),
3398 TapeOp::Tan(a) => vals[*a].tan(),
3399 TapeOp::Atan(a) => vals[*a].atan(),
3400 TapeOp::Acos(a) => vals[*a].acos(),
3401 TapeOp::Sinh(a) => vals[*a].sinh(),
3402 TapeOp::Cosh(a) => vals[*a].cosh(),
3403 TapeOp::Tanh(a) => vals[*a].tanh(),
3404 TapeOp::Asin(a) => vals[*a].asin(),
3405 TapeOp::Acosh(a) => vals[*a].acosh(),
3406 TapeOp::Asinh(a) => vals[*a].asinh(),
3407 TapeOp::Atanh(a) => vals[*a].atanh(),
3408 TapeOp::Erf(a) => erf(vals[*a]),
3409 TapeOp::XLogX(a) => xlogx(vals[*a]),
3410 TapeOp::CEntropy(a, b) => centropy(vals[*a], vals[*b]),
3411 TapeOp::Atan2(a, b) => vals[*a].atan2(vals[*b]),
3412 TapeOp::Cmp(_, _, _)
3413 | TapeOp::And(_, _)
3414 | TapeOp::Or(_, _)
3415 | TapeOp::Not(_)
3416 | TapeOp::Select(_, _, _)
3417 | TapeOp::Min(_, _)
3418 | TapeOp::Max(_, _) => panic!(
3419 "GlobalTape free-function kernels do not implement conditional / logical \
3420 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3421 instead."
3422 ),
3423 TapeOp::Funcall(fc) => {
3424 let FuncallData { lib, name, args } = fc.as_ref();
3425 let call_args = funcall_to_ext_args(args, vals);
3426 let res = lib
3427 .eval(name, &call_args, false, false)
3428 .unwrap_or_else(|e| panic!("external function '{name}' eval failed: {e}"));
3429 res.value
3430 }
3431 }
3432}
3433
3434#[inline]
3435fn rev_step(op: &TapeOp, i: usize, vals: &[f64], adj: &mut [f64], a: f64, grad: &mut [f64]) {
3436 match op {
3437 TapeOp::Const(_) => {}
3438 TapeOp::Var(j) => {
3439 grad[*j] += a;
3440 }
3441 TapeOp::Add(l, r) => {
3442 adj[*l] += a;
3443 adj[*r] += a;
3444 }
3445 TapeOp::Sub(l, r) => {
3446 adj[*l] += a;
3447 adj[*r] -= a;
3448 }
3449 TapeOp::Mul(l, r) => {
3450 adj[*l] += a * vals[*r];
3451 adj[*r] += a * vals[*l];
3452 }
3453 TapeOp::Div(l, r) => {
3454 let rv = vals[*r];
3456 adj[*l] += a / rv;
3457 adj[*r] -= a * vals[i] / rv;
3458 }
3459 TapeOp::Pow(l, r) => {
3460 let lv = vals[*l];
3461 let rv = vals[*r];
3462 if rv != 0.0 {
3463 adj[*l] += a * rv * lv.powf(rv - 1.0);
3464 }
3465 if lv > 0.0 {
3466 adj[*r] += a * vals[i] * lv.ln();
3467 }
3468 }
3469 TapeOp::Neg(j) => {
3470 adj[*j] -= a;
3471 }
3472 TapeOp::Abs(j) => {
3473 if vals[*j] >= 0.0 {
3474 adj[*j] += a;
3475 } else {
3476 adj[*j] -= a;
3477 }
3478 }
3479 TapeOp::Sqrt(j) => {
3480 let sv = vals[i];
3481 if sv > 0.0 {
3482 adj[*j] += a * 0.5 / sv;
3483 }
3484 }
3485 TapeOp::Exp(j) => {
3486 adj[*j] += a * vals[i];
3487 }
3488 TapeOp::Log(j) => {
3489 adj[*j] += a / vals[*j];
3490 }
3491 TapeOp::Log10(j) => {
3492 adj[*j] += a / (vals[*j] * std::f64::consts::LN_10);
3493 }
3494 TapeOp::Sin(j) => {
3495 adj[*j] += a * vals[*j].cos();
3496 }
3497 TapeOp::Cos(j) => {
3498 adj[*j] -= a * vals[*j].sin();
3499 }
3500 TapeOp::Tan(j) => {
3501 let t = vals[i];
3502 adj[*j] += a * (1.0 + t * t);
3503 }
3504 TapeOp::Atan(j) => {
3505 let u = vals[*j];
3506 adj[*j] += a / (1.0 + u * u);
3507 }
3508 TapeOp::Acos(j) => {
3509 let u = vals[*j];
3510 adj[*j] -= a / (1.0 - u * u).sqrt();
3511 }
3512 TapeOp::Sinh(j) => {
3513 adj[*j] += a * vals[*j].cosh();
3514 }
3515 TapeOp::Cosh(j) => {
3516 adj[*j] += a * vals[*j].sinh();
3517 }
3518 TapeOp::Tanh(j) => {
3519 let t = vals[i];
3520 adj[*j] += a * (1.0 - t * t);
3521 }
3522 TapeOp::Asin(j) => {
3523 let u = vals[*j];
3524 adj[*j] += a / (1.0 - u * u).sqrt();
3525 }
3526 TapeOp::Acosh(j) => {
3527 let u = vals[*j];
3528 adj[*j] += a / (u * u - 1.0).sqrt();
3529 }
3530 TapeOp::Asinh(j) => {
3531 let u = vals[*j];
3532 adj[*j] += a / (u * u + 1.0).sqrt();
3533 }
3534 TapeOp::Atanh(j) => {
3535 let u = vals[*j];
3536 adj[*j] += a / (1.0 - u * u);
3537 }
3538 TapeOp::Erf(j) => {
3539 adj[*j] += a * erf_d1(vals[*j]);
3540 }
3541 TapeOp::XLogX(j) => {
3542 adj[*j] += a * xlogx_d1(vals[*j]);
3543 }
3544 TapeOp::CEntropy(l, r) => {
3545 adj[*l] += a * centropy_da(vals[*l], vals[*r]);
3546 adj[*r] += a * centropy_db(vals[*l], vals[*r]);
3547 }
3548 TapeOp::Atan2(l, r) => {
3549 let y = vals[*l];
3550 let x = vals[*r];
3551 let d = y * y + x * x;
3552 adj[*l] += a * (x / d);
3553 adj[*r] += a * (-y / d);
3554 }
3555 TapeOp::Cmp(_, _, _)
3556 | TapeOp::And(_, _)
3557 | TapeOp::Or(_, _)
3558 | TapeOp::Not(_)
3559 | TapeOp::Select(_, _, _)
3560 | TapeOp::Min(_, _)
3561 | TapeOp::Max(_, _) => panic!(
3562 "GlobalTape free-function kernels do not implement conditional / logical \
3563 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3564 instead."
3565 ),
3566 TapeOp::Funcall(fc) => {
3567 let FuncallData { lib, name, args } = fc.as_ref();
3568 let call_args = funcall_to_ext_args(args, vals);
3569 let res = lib
3570 .eval(name, &call_args, true, false)
3571 .unwrap_or_else(|e| panic!("external function '{name}' reverse eval failed: {e}"));
3572 let derivs = res.derivs.expect("want_derivs=true returns derivs");
3573 let mut k = 0usize;
3574 for arg in args {
3575 if let TapeFuncallArg::Tape(idx) = arg {
3576 adj[*idx] += a * derivs[k];
3577 k += 1;
3578 }
3579 }
3580 let _ = i;
3581 let _ = grad;
3582 }
3583 }
3584}
3585
3586#[inline]
3594fn fwd_dir_step(op: &TapeOp, seed: &[f64], vals: &[f64], dot: &[f64], i: usize) -> f64 {
3595 match op {
3596 TapeOp::Const(_) => 0.0,
3597 TapeOp::Var(k) => seed[*k],
3598 TapeOp::Add(a, b) => dot[*a] + dot[*b],
3599 TapeOp::Sub(a, b) => dot[*a] - dot[*b],
3600 TapeOp::Mul(a, b) => dot[*a] * vals[*b] + vals[*a] * dot[*b],
3601 TapeOp::Div(a, b) => {
3602 (dot[*a] - vals[i] * dot[*b]) / vals[*b]
3605 }
3606 TapeOp::Pow(a, b) => {
3607 let u = vals[*a];
3608 let r = vals[*b];
3609 let du = dot[*a];
3610 let dr = dot[*b];
3611 let mut result = 0.0;
3612 if r != 0.0 {
3617 result += r * u.powf(r - 1.0) * du;
3618 }
3619 if u > 0.0 {
3620 result += vals[i] * u.ln() * dr;
3621 }
3622 result
3623 }
3624 TapeOp::Neg(a) => -dot[*a],
3625 TapeOp::Abs(a) => {
3626 if vals[*a] >= 0.0 {
3627 dot[*a]
3628 } else {
3629 -dot[*a]
3630 }
3631 }
3632 TapeOp::Sqrt(a) => {
3633 let sv = vals[i];
3634 if sv > 0.0 { dot[*a] * 0.5 / sv } else { 0.0 }
3635 }
3636 TapeOp::Exp(a) => vals[i] * dot[*a],
3637 TapeOp::Log(a) => dot[*a] / vals[*a],
3638 TapeOp::Log10(a) => dot[*a] / (vals[*a] * std::f64::consts::LN_10),
3639 TapeOp::Sin(a) => vals[*a].cos() * dot[*a],
3640 TapeOp::Cos(a) => -vals[*a].sin() * dot[*a],
3641 TapeOp::Tan(a) => {
3642 let t = vals[i];
3643 (1.0 + t * t) * dot[*a]
3644 }
3645 TapeOp::Atan(a) => {
3646 let u = vals[*a];
3647 dot[*a] / (1.0 + u * u)
3648 }
3649 TapeOp::Acos(a) => {
3650 let u = vals[*a];
3651 -dot[*a] / (1.0 - u * u).sqrt()
3652 }
3653 TapeOp::Sinh(a) => dot[*a] * vals[*a].cosh(),
3654 TapeOp::Cosh(a) => dot[*a] * vals[*a].sinh(),
3655 TapeOp::Tanh(a) => {
3656 let t = vals[i];
3657 (1.0 - t * t) * dot[*a]
3658 }
3659 TapeOp::Asin(a) => {
3660 let u = vals[*a];
3661 dot[*a] / (1.0 - u * u).sqrt()
3662 }
3663 TapeOp::Acosh(a) => {
3664 let u = vals[*a];
3665 dot[*a] / (u * u - 1.0).sqrt()
3666 }
3667 TapeOp::Asinh(a) => {
3668 let u = vals[*a];
3669 dot[*a] / (u * u + 1.0).sqrt()
3670 }
3671 TapeOp::Atanh(a) => {
3672 let u = vals[*a];
3673 dot[*a] / (1.0 - u * u)
3674 }
3675 TapeOp::Erf(a) => erf_d1(vals[*a]) * dot[*a],
3676 TapeOp::XLogX(a) => xlogx_d1(vals[*a]) * dot[*a],
3677 TapeOp::CEntropy(a, b) => {
3678 centropy_da(vals[*a], vals[*b]) * dot[*a] + centropy_db(vals[*a], vals[*b]) * dot[*b]
3679 }
3680 TapeOp::Atan2(a, b) => {
3681 let y = vals[*a];
3682 let x = vals[*b];
3683 let d = y * y + x * x;
3684 (x * dot[*a] - y * dot[*b]) / d
3685 }
3686 TapeOp::Cmp(_, _, _)
3687 | TapeOp::And(_, _)
3688 | TapeOp::Or(_, _)
3689 | TapeOp::Not(_)
3690 | TapeOp::Select(_, _, _)
3691 | TapeOp::Min(_, _)
3692 | TapeOp::Max(_, _) => panic!(
3693 "GlobalTape free-function kernels do not implement conditional / logical \
3694 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3695 instead."
3696 ),
3697 TapeOp::Funcall(fc) => {
3698 let FuncallData { lib, name, args } = fc.as_ref();
3699 let call_args = funcall_to_ext_args(args, vals);
3700 let res = lib
3701 .eval(name, &call_args, true, false)
3702 .unwrap_or_else(|e| panic!("external function '{name}' tangent eval failed: {e}"));
3703 let derivs = res.derivs.expect("want_derivs=true returns derivs");
3704 let mut acc = 0.0;
3705 let mut k = 0usize;
3706 for arg in args {
3707 if let TapeFuncallArg::Tape(idx) = arg {
3708 acc += derivs[k] * dot[*idx];
3709 k += 1;
3710 }
3711 }
3712 let _ = seed;
3713 acc
3714 }
3715 }
3716}
3717
3718#[allow(clippy::too_many_arguments)]
3727#[inline]
3728fn ror_dir_step(
3729 op: &TapeOp,
3730 i: usize,
3731 vals: &[f64],
3732 dot: &[f64],
3733 adj: &mut [f64],
3734 adj_dot: &mut [f64],
3735 w: f64,
3736 wd: f64,
3737 weight: f64,
3738 out: &mut [f64],
3739) {
3740 match op {
3741 TapeOp::Const(_) => {}
3742 TapeOp::Var(k) => {
3743 if wd != 0.0 {
3744 out[*k] += weight * wd;
3745 }
3746 }
3747 TapeOp::Add(a, b) => {
3748 adj[*a] += w;
3749 adj[*b] += w;
3750 adj_dot[*a] += wd;
3751 adj_dot[*b] += wd;
3752 }
3753 TapeOp::Sub(a, b) => {
3754 adj[*a] += w;
3755 adj[*b] -= w;
3756 adj_dot[*a] += wd;
3757 adj_dot[*b] -= wd;
3758 }
3759 TapeOp::Mul(a, b) => {
3760 adj[*a] += w * vals[*b];
3761 adj[*b] += w * vals[*a];
3762 adj_dot[*a] += wd * vals[*b] + w * dot[*b];
3763 adj_dot[*b] += wd * vals[*a] + w * dot[*a];
3764 }
3765 TapeOp::Div(a, b) => {
3766 let vb = vals[*b];
3769 let q = vals[i];
3770 let qd = dot[i];
3771 adj[*a] += w / vb;
3772 adj_dot[*a] += wd / vb - w * (dot[*b] / vb) / vb;
3773 adj[*b] -= w * q / vb;
3774 adj_dot[*b] += -(wd * q) / vb + (w / vb) * (-qd + q * (dot[*b] / vb));
3775 }
3776 TapeOp::Pow(a, b) => {
3777 let u = vals[*a];
3778 let r = vals[*b];
3779 let du = dot[*a];
3780 let dr = dot[*b];
3781 if r != 0.0 {
3782 if u != 0.0 {
3783 let p_a = r * u.powf(r - 1.0);
3784 adj[*a] += w * p_a;
3785 let mut dp_a = dr * u.powf(r - 1.0);
3786 if u > 0.0 {
3787 dp_a += r * u.powf(r - 1.0) * ((r - 1.0) * du / u + dr * u.ln());
3788 } else {
3789 dp_a += r * (r - 1.0) * u.powf(r - 2.0) * du;
3790 }
3791 adj_dot[*a] += wd * p_a + w * dp_a;
3792 } else if r >= 2.0 {
3793 let p_a = 0.0;
3794 adj[*a] += w * p_a;
3795 let dp_a = if r == 2.0 {
3796 2.0 * du
3797 } else {
3798 r * (r - 1.0) * (0.0_f64).powf(r - 2.0) * du
3799 };
3800 adj_dot[*a] += wd * p_a + w * dp_a;
3801 }
3802 }
3803 if u > 0.0 {
3804 let ln_u = u.ln();
3805 let p_b = vals[i] * ln_u;
3806 adj[*b] += w * p_b;
3807 let dur = vals[i] * (r * du / u + dr * ln_u);
3808 let dp_b = dur * ln_u + vals[i] * du / u;
3809 adj_dot[*b] += wd * p_b + w * dp_b;
3810 }
3811 }
3812 TapeOp::Neg(a) => {
3813 adj[*a] -= w;
3814 adj_dot[*a] -= wd;
3815 }
3816 TapeOp::Abs(a) => {
3817 let s = if vals[*a] >= 0.0 { 1.0 } else { -1.0 };
3818 adj[*a] += w * s;
3819 adj_dot[*a] += wd * s;
3820 }
3821 TapeOp::Sqrt(a) => {
3822 let sv = vals[i];
3823 if sv > 0.0 {
3824 let fp = 0.5 / sv;
3825 let fpp = -0.25 / (vals[*a] * sv);
3826 adj[*a] += w * fp;
3827 adj_dot[*a] += wd * fp + w * fpp * dot[*a];
3828 }
3829 }
3830 TapeOp::Exp(a) => {
3831 let ev = vals[i];
3832 adj[*a] += w * ev;
3833 adj_dot[*a] += wd * ev + w * ev * dot[*a];
3834 }
3835 TapeOp::Log(a) => {
3836 let u = vals[*a];
3837 adj[*a] += w / u;
3838 adj_dot[*a] += wd / u + w * (-1.0 / (u * u)) * dot[*a];
3839 }
3840 TapeOp::Log10(a) => {
3841 let u = vals[*a];
3842 let c = std::f64::consts::LN_10;
3843 adj[*a] += w / (u * c);
3844 adj_dot[*a] += wd / (u * c) + w * (-1.0 / (u * u * c)) * dot[*a];
3845 }
3846 TapeOp::Sin(a) => {
3847 let u = vals[*a];
3848 let cu = u.cos();
3849 adj[*a] += w * cu;
3850 adj_dot[*a] += wd * cu + w * (-u.sin()) * dot[*a];
3851 }
3852 TapeOp::Cos(a) => {
3853 let u = vals[*a];
3854 let su = u.sin();
3855 adj[*a] -= w * su;
3856 adj_dot[*a] += wd * (-su) + w * (-u.cos()) * dot[*a];
3857 }
3858 TapeOp::Tan(a) => {
3859 let t = vals[i];
3860 let gp = 1.0 + t * t;
3861 let gpp = 2.0 * t * gp;
3862 adj[*a] += w * gp;
3863 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3864 }
3865 TapeOp::Atan(a) => {
3866 let u = vals[*a];
3867 let d = 1.0 + u * u;
3868 let gp = 1.0 / d;
3869 let gpp = -2.0 * u / (d * d);
3870 adj[*a] += w * gp;
3871 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3872 }
3873 TapeOp::Acos(a) => {
3874 let u = vals[*a];
3875 let s = 1.0 - u * u;
3876 let r = s.sqrt();
3877 let gp = -1.0 / r;
3878 let gpp = -u / (s * r);
3879 adj[*a] += w * gp;
3880 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3881 }
3882 TapeOp::Sinh(a) => {
3883 let u = vals[*a];
3884 let gp = u.cosh();
3885 let gpp = vals[i]; adj[*a] += w * gp;
3887 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3888 }
3889 TapeOp::Cosh(a) => {
3890 let u = vals[*a];
3891 let gp = u.sinh();
3892 let gpp = vals[i]; adj[*a] += w * gp;
3894 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3895 }
3896 TapeOp::Tanh(a) => {
3897 let t = vals[i];
3898 let gp = 1.0 - t * t;
3899 let gpp = -2.0 * t * gp;
3900 adj[*a] += w * gp;
3901 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3902 }
3903 TapeOp::Asin(a) => {
3904 let u = vals[*a];
3905 let s = 1.0 - u * u;
3906 let r = s.sqrt();
3907 let gp = 1.0 / r;
3908 let gpp = u / (s * r);
3909 adj[*a] += w * gp;
3910 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3911 }
3912 TapeOp::Acosh(a) => {
3913 let u = vals[*a];
3914 let s = u * u - 1.0;
3915 let r = s.sqrt();
3916 let gp = 1.0 / r;
3917 let gpp = -u / (s * r);
3918 adj[*a] += w * gp;
3919 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3920 }
3921 TapeOp::Asinh(a) => {
3922 let u = vals[*a];
3923 let s = u * u + 1.0;
3924 let r = s.sqrt();
3925 let gp = 1.0 / r;
3926 let gpp = -u / (s * r);
3927 adj[*a] += w * gp;
3928 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3929 }
3930 TapeOp::Atanh(a) => {
3931 let u = vals[*a];
3932 let d = 1.0 - u * u;
3933 let gp = 1.0 / d;
3934 let gpp = 2.0 * u / (d * d);
3935 adj[*a] += w * gp;
3936 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3937 }
3938 TapeOp::Erf(a) => {
3939 let u = vals[*a];
3940 let gp = erf_d1(u);
3941 let gpp = erf_d2(u);
3942 adj[*a] += w * gp;
3943 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3944 }
3945 TapeOp::XLogX(a) => {
3946 let u = vals[*a];
3949 let gp = xlogx_d1(u);
3950 let gpp = xlogx_d2(u);
3951 adj[*a] += w * gp;
3952 adj_dot[*a] += wd * gp + w * gpp * dot[*a];
3953 }
3954 TapeOp::CEntropy(a, b) => {
3955 let ua = vals[*a];
3956 let ub = vals[*b];
3957 let fa = centropy_da(ua, ub);
3958 let fb = centropy_db(ua, ub);
3959 let faa = centropy_daa(ua);
3960 let fab = centropy_dab(ub);
3961 let fbb = centropy_dbb(ua, ub);
3962 adj[*a] += w * fa;
3963 adj[*b] += w * fb;
3964 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
3965 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
3966 }
3967 TapeOp::Atan2(a, b) => {
3968 let y = vals[*a];
3969 let x = vals[*b];
3970 let d = y * y + x * x;
3971 let d2 = d * d;
3972 let fa = x / d;
3973 let fb = -y / d;
3974 let faa = -2.0 * x * y / d2;
3975 let fab = (y * y - x * x) / d2;
3976 let fbb = 2.0 * x * y / d2;
3977 adj[*a] += w * fa;
3978 adj[*b] += w * fb;
3979 adj_dot[*a] += wd * fa + w * (faa * dot[*a] + fab * dot[*b]);
3980 adj_dot[*b] += wd * fb + w * (fab * dot[*a] + fbb * dot[*b]);
3981 }
3982 TapeOp::Cmp(_, _, _)
3983 | TapeOp::And(_, _)
3984 | TapeOp::Or(_, _)
3985 | TapeOp::Not(_)
3986 | TapeOp::Select(_, _, _)
3987 | TapeOp::Min(_, _)
3988 | TapeOp::Max(_, _) => panic!(
3989 "GlobalTape free-function kernels do not implement conditional / logical \
3990 / min-max TapeOps; use the Tape (build_with_externals) interpreter path \
3991 instead."
3992 ),
3993 TapeOp::Funcall(fc) => {
3994 let FuncallData { lib, name, args } = fc.as_ref();
3995 let call_args = funcall_to_ext_args(args, vals);
3996 let res = lib.eval(name, &call_args, true, true).unwrap_or_else(|e| {
3997 panic!("external function '{name}' 2nd-order eval failed: {e}")
3998 });
3999 let derivs = res.derivs.expect("want_derivs=true returns derivs");
4000 let hes = res.hessian.expect("want_hes=true returns hessian");
4001 let real_tape: Vec<usize> = args
4002 .iter()
4003 .filter_map(|a| match a {
4004 TapeFuncallArg::Tape(t) => Some(*t),
4005 TapeFuncallArg::Str(_) => None,
4006 })
4007 .collect();
4008 for (k, &tk) in real_tape.iter().enumerate() {
4009 adj[tk] += w * derivs[k];
4010 let mut second_term = 0.0;
4011 for (l, &tl) in real_tape.iter().enumerate() {
4012 let (lo, hi) = if k <= l { (k, l) } else { (l, k) };
4013 let h_kl = hes[lo + hi * (hi + 1) / 2];
4014 second_term += h_kl * dot[tl];
4015 }
4016 adj_dot[tk] += wd * derivs[k] + w * second_term;
4017 }
4018 let _ = out;
4019 let _ = weight;
4020 let _ = i;
4021 }
4022 }
4023}
4024
4025fn hessian_sparsity_impl(ops: &[TapeOp]) -> BTreeSet<(usize, usize)> {
4029 let n = ops.len();
4030 let mut var_sets: Vec<BTreeSet<usize>> = Vec::with_capacity(n);
4031 let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new();
4032
4033 let emit_cross =
4034 |s1: &BTreeSet<usize>, s2: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4035 for &v1 in s1 {
4036 for &v2 in s2 {
4037 let (r, c) = if v1 >= v2 { (v1, v2) } else { (v2, v1) };
4038 pairs.insert((r, c));
4039 }
4040 }
4041 };
4042 let emit_self = |s: &BTreeSet<usize>, pairs: &mut BTreeSet<(usize, usize)>| {
4043 let vars: Vec<usize> = s.iter().copied().collect();
4044 for (ai, &vi) in vars.iter().enumerate() {
4045 for &vj in &vars[..=ai] {
4046 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4047 pairs.insert((r, c));
4048 }
4049 }
4050 };
4051
4052 for op in ops {
4053 let vset = match op {
4054 TapeOp::Const(_) => BTreeSet::new(),
4055 TapeOp::Var(j) => {
4056 let mut s = BTreeSet::new();
4057 s.insert(*j);
4058 s
4059 }
4060 TapeOp::Add(a, b) | TapeOp::Sub(a, b) => {
4061 var_sets[*a].union(&var_sets[*b]).copied().collect()
4062 }
4063 TapeOp::Neg(a) | TapeOp::Abs(a) => var_sets[*a].clone(),
4064 TapeOp::Mul(a, b) => {
4065 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4066 var_sets[*a].union(&var_sets[*b]).copied().collect()
4067 }
4068 TapeOp::Div(a, b) => {
4069 emit_cross(&var_sets[*a], &var_sets[*b], &mut pairs);
4070 emit_self(&var_sets[*b], &mut pairs);
4071 var_sets[*a].union(&var_sets[*b]).copied().collect()
4072 }
4073 TapeOp::Pow(a, b) | TapeOp::Atan2(a, b) | TapeOp::CEntropy(a, b) => {
4074 let combined: BTreeSet<usize> =
4075 var_sets[*a].union(&var_sets[*b]).copied().collect();
4076 emit_self(&combined, &mut pairs);
4077 combined
4078 }
4079 TapeOp::Sqrt(a)
4080 | TapeOp::Exp(a)
4081 | TapeOp::Log(a)
4082 | TapeOp::Log10(a)
4083 | TapeOp::Sin(a)
4084 | TapeOp::Cos(a)
4085 | TapeOp::Tan(a)
4086 | TapeOp::Atan(a)
4087 | TapeOp::Acos(a)
4088 | TapeOp::Sinh(a)
4089 | TapeOp::Cosh(a)
4090 | TapeOp::Tanh(a)
4091 | TapeOp::Asin(a)
4092 | TapeOp::Acosh(a)
4093 | TapeOp::Asinh(a)
4094 | TapeOp::Erf(a)
4095 | TapeOp::XLogX(a)
4096 | TapeOp::Atanh(a) => {
4097 emit_self(&var_sets[*a], &mut pairs);
4098 var_sets[*a].clone()
4099 }
4100 TapeOp::Funcall(fc) => {
4101 let args = &fc.args;
4102 let mut combined: BTreeSet<usize> = BTreeSet::new();
4103 for arg in args {
4104 if let TapeFuncallArg::Tape(t) = arg {
4105 for &vv in &var_sets[*t] {
4106 combined.insert(vv);
4107 }
4108 }
4109 }
4110 emit_self(&combined, &mut pairs);
4111 combined
4112 }
4113 TapeOp::Cmp(_, _, _) | TapeOp::And(_, _) | TapeOp::Or(_, _) | TapeOp::Not(_) => {
4114 BTreeSet::new()
4117 }
4118 TapeOp::Select(_, t, e) => {
4119 var_sets[*t].union(&var_sets[*e]).copied().collect()
4122 }
4123 TapeOp::Min(a, b) | TapeOp::Max(a, b) => {
4124 var_sets[*a].union(&var_sets[*b]).copied().collect()
4127 }
4128 };
4129 var_sets.push(vset);
4130 }
4131 pairs
4132}
4133
4134#[cfg(test)]
4135mod tests {
4136 use super::*;
4137
4138 fn cnst(c: f64) -> Expr {
4139 Expr::Const(c)
4140 }
4141 fn var(i: usize) -> Expr {
4142 Expr::Var(i)
4143 }
4144 fn add(a: Expr, b: Expr) -> Expr {
4145 Expr::Binary(BinOp::Add, Box::new(a), Box::new(b))
4146 }
4147 fn mul(a: Expr, b: Expr) -> Expr {
4148 Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b))
4149 }
4150 fn pow(a: Expr, b: Expr) -> Expr {
4151 Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b))
4152 }
4153 fn div(a: Expr, b: Expr) -> Expr {
4154 Expr::Binary(BinOp::Div, Box::new(a), Box::new(b))
4155 }
4156 fn unary(op: UnaryOp, a: Expr) -> Expr {
4157 Expr::Unary(op, Box::new(a))
4158 }
4159 fn cmp(op: CmpOp, a: Expr, b: Expr) -> Expr {
4160 Expr::Compare(op, Box::new(a), Box::new(b))
4161 }
4162 fn cond(c: Expr, t: Expr, e: Expr) -> Expr {
4163 Expr::Cond {
4164 cond: Box::new(c),
4165 then_: Box::new(t),
4166 else_: Box::new(e),
4167 }
4168 }
4169
4170 #[test]
4171 fn polynomial_eval_and_grad() {
4172 let e = add(
4174 mul(cnst(3.0), pow(var(0), cnst(2.0))),
4175 mul(cnst(2.0), var(1)),
4176 );
4177 let t = Tape::build(&e);
4178 assert!((t.eval(&[2.0, 3.0]) - 18.0).abs() < 1e-12);
4179 let mut g = vec![0.0; 2];
4180 t.gradient_seed(&[2.0, 3.0], 1.0, &mut g);
4181 assert!((g[0] - 12.0).abs() < 1e-12);
4183 assert!((g[1] - 2.0).abs() < 1e-12);
4184 }
4185
4186 #[test]
4187 fn cse_shared_body_evaluated_once() {
4188 let body = Arc::new(add(var(0), var(1)));
4190 let e = add(
4191 pow(Expr::Cse(body.clone()), cnst(2.0)),
4192 Expr::Cse(body.clone()),
4193 );
4194 let t = Tape::build(&e);
4195 let n_body_adds = t
4197 .ops
4198 .iter()
4199 .filter(|op| {
4200 matches!(op, TapeOp::Add(a, b) if {
4201 matches!(t.ops[*a], TapeOp::Var(0)) && matches!(t.ops[*b], TapeOp::Var(1))
4202 })
4203 })
4204 .count();
4205 assert_eq!(n_body_adds, 1, "CSE body should be emitted exactly once");
4206
4207 assert!((t.eval(&[1.0, 2.0]) - 12.0).abs() < 1e-12);
4209 let mut g = vec![0.0; 2];
4210 t.gradient_seed(&[1.0, 2.0], 1.0, &mut g);
4211 assert!((g[0] - 7.0).abs() < 1e-12);
4213 assert!((g[1] - 7.0).abs() < 1e-12);
4214 }
4215
4216 fn fd_check(tape: &Tape, x: &[f64], n: usize, tol: f64) {
4217 let vars = tape.variables();
4218 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4219 let mut pairs = Vec::new();
4220 for (ai, &vi) in vars.iter().enumerate() {
4221 for &vj in &vars[..=ai] {
4222 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4223 hess_map.entry((r, c)).or_insert_with(|| {
4224 let p = pairs.len();
4225 pairs.push((r, c));
4226 p
4227 });
4228 }
4229 }
4230 let nnz = pairs.len();
4231 let mut ad = vec![0.0; nnz];
4232 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
4233
4234 let mut fd = vec![0.0; nnz];
4235 let mut xp = x.to_vec();
4236 let mut gp = vec![0.0; n];
4237 let mut gm = vec![0.0; n];
4238 for &j in &vars {
4239 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4240 xp[j] = x[j] + h;
4241 gp.iter_mut().for_each(|v| *v = 0.0);
4242 tape.gradient_seed(&xp, 1.0, &mut gp);
4243 xp[j] = x[j] - h;
4244 gm.iter_mut().for_each(|v| *v = 0.0);
4245 tape.gradient_seed(&xp, 1.0, &mut gm);
4246 xp[j] = x[j];
4247 for &i in &vars {
4248 if i >= j {
4249 if let Some(&pos) = hess_map.get(&(i, j)) {
4250 fd[pos] = (gp[i] - gm[i]) / (2.0 * h);
4251 }
4252 }
4253 }
4254 }
4255 for (k, &(r, c)) in pairs.iter().enumerate() {
4256 let scale = fd[k].abs().max(1.0);
4257 assert!(
4258 (ad[k] - fd[k]).abs() / scale < tol,
4259 "H[{},{}]: AD={:.6e} FD={:.6e}",
4260 r,
4261 c,
4262 ad[k],
4263 fd[k]
4264 );
4265 }
4266 }
4267
4268 #[test]
4269 fn hessian_quadratic_matches_fd() {
4270 let e = add(
4272 add(
4273 mul(cnst(3.0), pow(var(0), cnst(2.0))),
4274 mul(cnst(2.0), mul(var(0), var(1))),
4275 ),
4276 pow(var(1), cnst(2.0)),
4277 );
4278 let t = Tape::build(&e);
4279 fd_check(&t, &[2.0, 3.0], 2, 1e-5);
4280 }
4281
4282 #[test]
4283 fn hessian_transcendental_matches_fd() {
4284 let e = Expr::Sum(vec![
4286 unary(UnaryOp::Exp, var(0)),
4287 unary(UnaryOp::Sin, var(1)),
4288 unary(UnaryOp::Log, var(0)),
4289 unary(UnaryOp::Sqrt, var(1)),
4290 mul(var(0), var(1)),
4291 ]);
4292 let t = Tape::build(&e);
4293 fd_check(&t, &[1.5, 2.0], 2, 1e-5);
4294 }
4295
4296 #[test]
4297 fn inverse_trig_grad_and_hessian_match_fd() {
4298 let e = Expr::Sum(vec![
4302 unary(UnaryOp::Tan, var(0)),
4303 unary(UnaryOp::Atan, var(1)),
4304 unary(UnaryOp::Acos, var(2)),
4305 mul(var(0), var(1)),
4306 ]);
4307 let t = Tape::build(&e);
4308 let x = [0.5, 1.3, 0.3];
4309
4310 let mut g = vec![0.0; 3];
4314 t.gradient_seed(&x, 1.0, &mut g);
4315 for j in 0..3 {
4316 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4317 let mut xp = x;
4318 let mut xm = x;
4319 xp[j] += h;
4320 xm[j] -= h;
4321 let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4322 let scale = fd.abs().max(1.0);
4323 assert!(
4324 (g[j] - fd).abs() / scale < 1e-5,
4325 "grad[{j}]: AD={:.6e} FD={:.6e}",
4326 g[j],
4327 fd
4328 );
4329 }
4330
4331 fd_check(&t, &x, 3, 1e-5);
4333 }
4334
4335 fn grad_and_hess_match_fd(e: &Expr, x: &[f64], tol: f64) {
4338 let n = x.len();
4339 let t = Tape::build(e);
4340 let mut g = vec![0.0; n];
4341 t.gradient_seed(x, 1.0, &mut g);
4342 for j in 0..n {
4343 let h = (1e-7_f64).max(x[j].abs() * 1e-7);
4344 let mut xp = x.to_vec();
4345 let mut xm = x.to_vec();
4346 xp[j] += h;
4347 xm[j] -= h;
4348 let fd = (t.eval(&xp) - t.eval(&xm)) / (2.0 * h);
4349 let scale = fd.abs().max(1.0);
4350 assert!(
4351 (g[j] - fd).abs() / scale < tol,
4352 "grad[{j}]: AD={:.6e} FD={:.6e}",
4353 g[j],
4354 fd
4355 );
4356 }
4357 fd_check(&t, x, n, tol);
4358 }
4359
4360 #[test]
4361 fn hyperbolic_grad_and_hessian_match_fd() {
4362 let e = Expr::Sum(vec![
4365 unary(UnaryOp::Sinh, var(0)),
4366 unary(UnaryOp::Cosh, var(1)),
4367 unary(UnaryOp::Tanh, var(2)),
4368 unary(UnaryOp::Asinh, var(3)),
4369 mul(var(0), var(1)),
4370 mul(var(2), var(3)),
4371 ]);
4372 grad_and_hess_match_fd(&e, &[0.5, 0.7, 0.3, 1.1], 1e-5);
4373 }
4374
4375 #[test]
4376 fn restricted_inverse_grad_and_hessian_match_fd() {
4377 let e = Expr::Sum(vec![
4381 unary(UnaryOp::Asin, var(0)),
4382 unary(UnaryOp::Acosh, var(1)),
4383 unary(UnaryOp::Atanh, var(2)),
4384 mul(var(0), var(2)),
4385 ]);
4386 grad_and_hess_match_fd(&e, &[0.4, 1.8, 0.3], 1e-5);
4387 }
4388
4389 #[test]
4390 fn erf_value_matches_reference() {
4391 let t = Tape::build(&unary(UnaryOp::Erf, var(0)));
4396 for (x, want) in [
4397 (0.0, 0.0),
4398 (0.5, 0.520_499_877_813_046_5),
4399 (1.0, 0.842_700_792_949_714_9),
4400 (-1.0, -0.842_700_792_949_714_9),
4401 (2.0, 0.995_322_265_018_952_7),
4402 (3.0, 0.999_977_909_503_001_4),
4403 ] {
4404 let got = t.eval(&[x]);
4405 assert!(
4406 (got - want).abs() < 1e-15,
4407 "erf({x}): got {got:.17e}, want {want:.17e}"
4408 );
4409 }
4410 assert!((t.eval(&[-0.3]) + t.eval(&[0.3])).abs() < 1e-16);
4413 assert!((t.eval(&[10.0]) - 1.0).abs() < 1e-15);
4414 }
4415
4416 #[test]
4417 fn erf_second_derivative_stays_finite_at_extreme_magnitudes() {
4418 for u in [1e19, 1e150, 1e300, f64::MAX, f64::MAX / 2.0] {
4426 for signed in [u, -u] {
4427 let d2 = erf_d2(signed);
4428 assert!(
4429 d2.is_finite(),
4430 "erf_d2({signed:e}) = {d2} — must be finite (the limit is 0)"
4431 );
4432 }
4433 }
4434 let u = 0.7;
4436 let want = -2.0 * u * (2.0 / std::f64::consts::PI.sqrt()) * (-u * u).exp();
4437 assert!((erf_d2(u) - want).abs() < 1e-15, "{}", erf_d2(u));
4438 }
4439
4440 #[test]
4441 fn quotient_rule_keeps_representable_derivatives_at_extreme_denominators() {
4442 let t = Tape::build(&div(var(0), var(1)));
4453 let mut g = vec![0.0; 2];
4454 for b in [1e300_f64, 1e200, 1e160, -1e300, 1e-160, 1e-300] {
4455 g.iter_mut().for_each(|v| *v = 0.0);
4461 t.gradient_seed(&[b, b], 1.0, &mut g);
4462 assert_eq!(g[0], 1.0 / b, "d(a/b)/da at a=b={b:e}");
4463 assert_eq!(
4464 g[1],
4465 -1.0 / b,
4466 "d(a/b)/db at a=b={b:e} — b² is not representable"
4467 );
4468 }
4469
4470 let b = 1e155_f64;
4474 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4475 hess_map.insert((0, 0), 0);
4476 hess_map.insert((1, 0), 1);
4477 hess_map.insert((1, 1), 2);
4478 let mut ad = vec![0.0; 3];
4479 t.hessian_accumulate(&[b, b], 1.0, &hess_map, &mut ad);
4480 let want = -(1.0 / b) / b;
4484 assert!(
4485 want != 0.0 && want.is_finite(),
4486 "test premise: -1/b² must be representable, got {want:e}"
4487 );
4488 assert!(
4489 (ad[1] - want).abs() <= 1e-10 * want.abs(),
4490 "d²(a/b)/da db at a=b={b:e}: got {:e}, want {want:e}",
4491 ad[1]
4492 );
4493 }
4494
4495 #[test]
4496 fn erf_grad_and_hessian_match_fd() {
4497 let e = Expr::Sum(vec![
4501 unary(UnaryOp::Erf, var(0)),
4502 unary(UnaryOp::Erf, mul(cnst(2.0), var(1))),
4503 mul(var(0), var(1)),
4504 ]);
4505 grad_and_hess_match_fd(&e, &[0.4, -0.7], 1e-5);
4506 }
4507
4508 #[test]
4509 fn erf_directional_hessian_matches_accumulated() {
4510 let e = Expr::Sum(vec![
4515 unary(UnaryOp::Erf, var(0)),
4516 unary(UnaryOp::Erf, mul(var(0), var(1))),
4517 ]);
4518 let tape = Tape::build(&e);
4519 let x = [0.6, -0.9];
4520 let n = x.len();
4521
4522 let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
4523 let hess_map: HashMap<(usize, usize), usize> =
4524 pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
4525 let mut acc = vec![0.0; pairs.len()];
4526 tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
4527
4528 let ops = tape.ops.len();
4530 let mut vals = vec![0.0; ops];
4531 tape.forward_into(&x, &mut vals);
4532 for j in 0..n {
4533 let mut seed = vec![0.0; n];
4534 seed[j] = 1.0;
4535 let mut col = vec![0.0; n];
4536 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
4537 tape.hessian_directional(
4538 &vals,
4539 &seed,
4540 1.0,
4541 &mut col,
4542 &mut dot,
4543 &mut adj,
4544 &mut adj_dot,
4545 );
4546 for i in 0..n {
4547 let (r, c) = if i >= j { (i, j) } else { (j, i) };
4548 let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
4549 assert!(
4550 (col[i] - want).abs() < 1e-12,
4551 "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
4552 col[i]
4553 );
4554 }
4555 }
4556 }
4557
4558 fn centropy_expr(a: Expr, b: Expr) -> Expr {
4559 Expr::Binary(BinOp::CEntropy, Box::new(a), Box::new(b))
4560 }
4561
4562 fn hess2(e: &Expr, x: &[f64; 2]) -> (f64, f64, f64) {
4564 let tape = Tape::build(e);
4565 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4566 hess_map.insert((0, 0), 0);
4567 hess_map.insert((1, 0), 1);
4568 hess_map.insert((1, 1), 2);
4569 let mut ad = vec![0.0; 3];
4570 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
4571 (ad[0], ad[1], ad[2])
4572 }
4573
4574 #[test]
4575 fn xlogx_and_centropy_grad_and_hessian_match_fd() {
4576 let e = Expr::Sum(vec![
4580 unary(UnaryOp::XLogX, var(0)),
4581 centropy_expr(var(0), var(1)),
4582 mul(var(0), var(1)),
4583 ]);
4584 grad_and_hess_match_fd(&e, &[1.7, 0.6], 1e-5);
4585 }
4586
4587 #[test]
4588 fn centropy_directional_hessian_matches_accumulated() {
4589 let e = Expr::Sum(vec![
4593 centropy_expr(var(0), var(1)),
4594 unary(UnaryOp::XLogX, mul(var(0), var(1))),
4595 ]);
4596 let tape = Tape::build(&e);
4597 let x = [1.3, 0.8];
4598 let n = x.len();
4599
4600 let pairs: Vec<(usize, usize)> = tape.hessian_sparsity().into_iter().collect();
4601 let hess_map: HashMap<(usize, usize), usize> =
4602 pairs.iter().enumerate().map(|(k, p)| (*p, k)).collect();
4603 let mut acc = vec![0.0; pairs.len()];
4604 tape.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
4605
4606 let ops = tape.ops.len();
4607 let mut vals = vec![0.0; ops];
4608 tape.forward_into(&x, &mut vals);
4609 let mut compared = 0usize;
4610 for j in 0..n {
4611 let mut seed = vec![0.0; n];
4612 seed[j] = 1.0;
4613 let mut col = vec![0.0; n];
4614 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
4615 tape.hessian_directional(
4616 &vals,
4617 &seed,
4618 1.0,
4619 &mut col,
4620 &mut dot,
4621 &mut adj,
4622 &mut adj_dot,
4623 );
4624 for i in 0..n {
4625 let (r, c) = if i >= j { (i, j) } else { (j, i) };
4626 let want = hess_map.get(&(r, c)).map_or(0.0, |&k| acc[k]);
4627 assert!(
4628 (col[i] - want).abs() < 1e-12,
4629 "H[{i},{j}]: directional={:.6e} accumulated={want:.6e}",
4630 col[i]
4631 );
4632 compared += 1;
4633 }
4634 }
4635 assert_eq!(compared, 4, "expected a full 2x2 comparison");
4638 assert!(
4639 acc.iter().any(|v| v.abs() > 1e-12),
4640 "every accumulated Hessian entry was zero — the arms under test never ran"
4641 );
4642 }
4643
4644 #[test]
4645 fn fused_entropy_ops_reach_derivatives_the_chain_rule_cannot() {
4646 let a = 1e-299_f64;
4652 let want = 1.0 / a; assert!(want.is_finite(), "test premise: 1/a must be representable");
4654
4655 let (fused, _, _) = hess2(&unary(UnaryOp::XLogX, var(0)), &[a, 0.0]);
4656 assert!(
4657 (fused - want).abs() <= 1e-12 * want,
4658 "xlogx''({a:e}): got {fused:e}, want {want:e}"
4659 );
4660
4661 let (decomposed, _, _) = hess2(&mul(var(0), unary(UnaryOp::Log, var(0))), &[a, 0.0]);
4666 assert!(
4667 !decomposed.is_finite(),
4668 "test premise: a*log(a) is supposed to lose this second derivative, \
4669 but it returned {decomposed:e} — if the Log/Mul arms now reach it, \
4670 re-derive whether the fused op is still needed"
4671 );
4672
4673 let (fused_aa, _, _) = hess2(¢ropy_expr(var(0), var(1)), &[a, 1.0]);
4675 assert!(
4676 (fused_aa - want).abs() <= 1e-12 * want,
4677 "centropy ∂²/∂a² at a={a:e}: got {fused_aa:e}, want {want:e}"
4678 );
4679 }
4680
4681 #[test]
4682 fn centropy_second_derivative_in_b_survives_a_squared_denominator() {
4683 let (a, b) = (1e300_f64, 1e200_f64);
4686 let want = (a / b) / b;
4687 assert!(
4691 want.is_finite() && want != 0.0 && a / (b * b) == 0.0,
4692 "test premise: a/b² representable ({want:e}) but a/(b*b) collapses"
4693 );
4694 let (_, _, h11) = hess2(¢ropy_expr(var(0), var(1)), &[a, b]);
4695 assert!(
4696 (h11 - want).abs() <= 1e-12 * want,
4697 "centropy ∂²/∂b² at ({a:e}, {b:e}): got {h11:e}, want {want:e}"
4698 );
4699 }
4700
4701 #[test]
4702 fn centropy_value_survives_an_out_of_range_ratio_and_the_zero_limit() {
4703 let t = |e: &Expr, x: &[f64]| Tape::build(e).eval(x);
4704 let ce = centropy_expr(var(0), var(1));
4705
4706 let want = 1e300 * (1e300_f64.ln() - 1e-300_f64.ln());
4710 let got = t(&ce, &[1e300, 1e-300]);
4711 assert!(
4712 got.is_finite() && (got - want).abs() <= 1e-12 * want,
4713 "centropy(1e300, 1e-300): got {got:e}, want {want:e}"
4714 );
4715 assert!(
4717 !t(
4718 &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
4719 &[1e300, 1e-300]
4720 )
4721 .is_finite(),
4722 "test premise: a*log(a/b) is supposed to overflow here"
4723 );
4724
4725 let (a, b) = (1e10 + 1.0, 1e10_f64);
4732 assert_eq!(a - b, 1.0, "test premise: a - b must be exact here");
4733 let got = t(&ce, &[a, b]);
4734 assert!(
4735 (got - 1.00000000005).abs() <= 2.3e-16,
4736 "centropy(1e10+1, 1e10): got {got:.17e}, want 1.00000000005"
4737 );
4738 let naive = t(
4740 &mul(var(0), unary(UnaryOp::Log, div(var(0), var(1)))),
4741 &[a, b],
4742 );
4743 assert!(
4744 (naive - 1.00000000005).abs() > 1e-9,
4745 "test premise: a*log(a/b) is supposed to lose digits here, got {naive:.17e}"
4746 );
4747
4748 assert_eq!(t(&ce, &[0.0, 2.0]), 0.0);
4750 assert_eq!(t(&unary(UnaryOp::XLogX, var(0)), &[0.0]), 0.0);
4751 }
4752
4753 #[test]
4754 fn cond_does_not_leak_a_non_finite_from_its_inactive_branch() {
4755 let cond = Expr::Compare(CmpOp::Lt, Box::new(var(0)), Box::new(cnst(1e-300)));
4763 let e = Expr::Cond {
4764 cond: Box::new(cond),
4765 then_: Box::new(mul(var(0), cnst(1e-300f64.ln()))),
4766 else_: Box::new(unary(UnaryOp::XLogX, var(0))),
4767 };
4768 let t = Tape::build(&e);
4769 let x = [-1.0_f64];
4770
4771 assert!(
4774 Tape::build(&unary(UnaryOp::XLogX, var(0)))
4775 .eval(&x)
4776 .is_nan(),
4777 "test premise: xlogx(-1) must be NaN"
4778 );
4779
4780 assert_eq!(t.eval(&x), -1e-300_f64.ln());
4781
4782 let mut g = vec![0.0; 1];
4783 t.gradient_seed(&x, 1.0, &mut g);
4784 assert_eq!(g[0], 1e-300_f64.ln(), "gradient leaked the inactive branch");
4785
4786 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4787 hess_map.insert((0, 0), 0);
4788 let mut acc = vec![0.0; 1];
4789 t.hessian_accumulate(&x, 1.0, &hess_map, &mut acc);
4790 assert_eq!(acc[0], 0.0, "hessian_accumulate leaked the inactive branch");
4791
4792 let ops = t.ops.len();
4793 let mut vals = vec![0.0; ops];
4794 t.forward_into(&x, &mut vals);
4795 let mut col = vec![0.0; 1];
4796 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; ops], vec![0.0; ops], vec![0.0; ops]);
4797 t.hessian_directional(
4798 &vals,
4799 &[1.0],
4800 1.0,
4801 &mut col,
4802 &mut dot,
4803 &mut adj,
4804 &mut adj_dot,
4805 );
4806 assert_eq!(
4807 col[0], 0.0,
4808 "hessian_directional leaked the inactive branch"
4809 );
4810 }
4811
4812 #[test]
4813 fn atan2_grad_and_hessian_match_fd() {
4814 let atan2 = |a: Expr, b: Expr| Expr::Binary(BinOp::Atan2, Box::new(a), Box::new(b));
4816 let e = Expr::Sum(vec![atan2(var(0), var(1)), mul(var(0), var(1))]);
4817 grad_and_hess_match_fd(&e, &[1.2, 0.7], 1e-5);
4818 }
4819
4820 #[test]
4821 fn minmax_grad_and_hessian_match_fd() {
4822 let e = Expr::Sum(vec![
4829 Expr::MinList(vec![var(0), var(1), var(2)]),
4830 Expr::MaxList(vec![var(1), var(2)]),
4831 mul(var(0), var(2)),
4832 ]);
4833 grad_and_hess_match_fd(&e, &[0.5, 3.0, 2.0], 1e-5);
4834 }
4835
4836 #[test]
4837 fn minmax_value_and_active_operand() {
4838 let e = Expr::Sum(vec![
4841 Expr::MinList(vec![var(0), var(1)]),
4842 Expr::MaxList(vec![var(0), var(1)]),
4843 ]);
4844 let t = Tape::build(&e);
4845 let x = [1.3, -0.4];
4847 assert!((t.eval(&x) - (x[0] + x[1])).abs() < 1e-12);
4848 let mut g = vec![0.0; 2];
4849 t.gradient_seed(&x, 1.0, &mut g);
4850 assert!((g[0] - 1.0).abs() < 1e-12, "g0={}", g[0]);
4853 assert!((g[1] - 1.0).abs() < 1e-12, "g1={}", g[1]);
4854 }
4855
4856 #[test]
4857 fn hessian_division_matches_fd() {
4858 let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
4860 let t = Tape::build(&e);
4861 fd_check(&t, &[0.5, 1.2], 2, 1e-5);
4862 }
4863
4864 #[test]
4865 fn conditional_value_grad_hessian_active_branch() {
4866 let e = cond(
4870 cmp(CmpOp::Ge, var(0), cnst(1.0)),
4871 mul(var(0), var(1)),
4872 pow(var(1), cnst(2.0)),
4873 );
4874 let t = Tape::build(&e);
4875
4876 let x = [2.0, 5.0];
4878 assert!((t.eval(&x) - 10.0).abs() < 1e-12);
4879 let mut g = vec![0.0; 2];
4880 t.gradient_seed(&x, 1.0, &mut g);
4881 assert!((g[0] - 5.0).abs() < 1e-10);
4883 assert!((g[1] - 2.0).abs() < 1e-10);
4884 fd_check(&t, &x, 2, 1e-5);
4886
4887 let x2 = [0.0, 5.0];
4889 assert!((t.eval(&x2) - 25.0).abs() < 1e-12);
4890 let mut g2 = vec![0.0; 2];
4891 t.gradient_seed(&x2, 1.0, &mut g2);
4892 assert!(g2[0].abs() < 1e-10);
4893 assert!((g2[1] - 10.0).abs() < 1e-10);
4894 fd_check(&t, &x2, 2, 1e-5);
4895 }
4896
4897 #[test]
4898 fn comparison_and_logical_have_zero_derivative() {
4899 let lt = cmp(CmpOp::Lt, var(0), var(1));
4903 let and = Expr::And(
4904 Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
4905 Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
4906 );
4907 let notc = Expr::Not(Box::new(cmp(CmpOp::Eq, var(0), var(1))));
4908 let e = add(add(lt, and), notc);
4909 let t = Tape::build(&e);
4910
4911 let x = [1.0, 2.0];
4912 assert!((t.eval(&x) - 3.0).abs() < 1e-12);
4914 let mut g = vec![0.0; 2];
4915 t.gradient_seed(&x, 1.0, &mut g);
4916 assert!(g[0].abs() < 1e-12, "d/dx0 should be 0, got {}", g[0]);
4917 assert!(g[1].abs() < 1e-12, "d/dx1 should be 0, got {}", g[1]);
4918 }
4919
4920 #[test]
4921 fn logical_or_value() {
4922 let e = Expr::Or(
4924 Box::new(cmp(CmpOp::Gt, var(0), cnst(0.0))),
4925 Box::new(cmp(CmpOp::Gt, var(1), cnst(0.0))),
4926 );
4927 let t = Tape::build(&e);
4928 assert!((t.eval(&[-1.0, 3.0]) - 1.0).abs() < 1e-12);
4929 assert!((t.eval(&[-1.0, -3.0]) - 0.0).abs() < 1e-12);
4930 }
4931
4932 fn directional_matches_accumulate(tape: &Tape, x: &[f64], n: usize) {
4937 let vars = tape.variables();
4938 let mut hess_map: HashMap<(usize, usize), usize> = HashMap::new();
4939 let mut pairs = Vec::new();
4940 for (ai, &vi) in vars.iter().enumerate() {
4941 for &vj in &vars[..=ai] {
4942 let (r, c) = if vi >= vj { (vi, vj) } else { (vj, vi) };
4943 hess_map.entry((r, c)).or_insert_with(|| {
4944 let p = pairs.len();
4945 pairs.push((r, c));
4946 p
4947 });
4948 }
4949 }
4950 let nnz = pairs.len();
4951 let mut ad = vec![0.0; nnz];
4952 tape.hessian_accumulate(x, 1.0, &hess_map, &mut ad);
4953
4954 let nops = tape.ops.len();
4955 let mut vals = vec![0.0; nops];
4956 tape.forward_into(x, &mut vals);
4957 let mut dot = vec![0.0; nops];
4958 let mut adj = vec![0.0; nops];
4959 let mut adj_dot = vec![0.0; nops];
4960
4961 for &j in &vars {
4962 let mut seed = vec![0.0; n];
4963 seed[j] = 1.0;
4964 let mut col = vec![0.0; n];
4965 tape.hessian_directional(
4966 &vals,
4967 &seed,
4968 1.0,
4969 &mut col,
4970 &mut dot,
4971 &mut adj,
4972 &mut adj_dot,
4973 );
4974 for &i in &vars {
4975 let (r, c) = if i >= j { (i, j) } else { (j, i) };
4976 let expect = ad[hess_map[&(r, c)]];
4977 assert!(
4978 (col[i] - expect).abs() < 1e-10,
4979 "directional H[{i},{j}] = {} vs accumulate {}",
4980 col[i],
4981 expect
4982 );
4983 }
4984 }
4985 }
4986
4987 #[test]
4988 fn directional_quadratic_matches_accumulate() {
4989 let e = add(
4991 add(
4992 mul(cnst(3.0), pow(var(0), cnst(2.0))),
4993 mul(mul(cnst(2.0), var(0)), var(1)),
4994 ),
4995 pow(var(1), cnst(2.0)),
4996 );
4997 let t = Tape::build(&e);
4998 directional_matches_accumulate(&t, &[0.5, -0.3], 2);
4999 }
5000
5001 #[test]
5002 fn directional_transcendental_matches_accumulate() {
5003 let e = Expr::Sum(vec![
5004 unary(UnaryOp::Exp, var(0)),
5005 unary(UnaryOp::Sin, var(1)),
5006 unary(UnaryOp::Log, var(0)),
5007 unary(UnaryOp::Sqrt, var(1)),
5008 mul(var(0), var(1)),
5009 ]);
5010 let t = Tape::build(&e);
5011 directional_matches_accumulate(&t, &[1.5, 2.0], 2);
5012 }
5013
5014 #[test]
5015 fn directional_with_division_matches_accumulate() {
5016 let e = add(div(var(0), var(1)), unary(UnaryOp::Cos, var(0)));
5017 let t = Tape::build(&e);
5018 directional_matches_accumulate(&t, &[0.5, 1.2], 2);
5019 }
5020
5021 #[test]
5022 fn hessian_sparsity_separable() {
5023 let e = add(unary(UnaryOp::Sin, var(0)), mul(var(1), var(2)));
5025 let t = Tape::build(&e);
5026 let s = t.hessian_sparsity();
5027 assert!(s.contains(&(0, 0)));
5028 assert!(s.contains(&(2, 1)));
5029 assert!(!s.contains(&(1, 0)));
5030 assert!(!s.contains(&(2, 0)));
5031 }
5032
5033 fn count_op<F: Fn(&TapeOp) -> bool>(t: &Tape, pred: F) -> usize {
5034 t.ops.iter().filter(|o| pred(o)).count()
5035 }
5036
5037 #[test]
5038 fn pow_zero_const_folds_to_one() {
5039 let e = pow(var(0), cnst(0.0));
5041 let t = Tape::build(&e);
5042 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5043 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Var(_))), 0);
5044 assert!((t.eval(&[7.0]) - 1.0).abs() < 1e-12);
5045 }
5046
5047 #[test]
5048 fn pow_one_passes_through() {
5049 let e = pow(var(0), cnst(1.0));
5051 let t = Tape::build(&e);
5052 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5053 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Const(_))), 0);
5054 assert!((t.eval(&[3.5]) - 3.5).abs() < 1e-12);
5055 }
5056
5057 #[test]
5058 fn pow_half_lowers_to_sqrt() {
5059 let e = pow(var(0), cnst(0.5));
5060 let t = Tape::build(&e);
5061 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5062 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Sqrt(_))), 1);
5063 assert!((t.eval(&[16.0]) - 4.0).abs() < 1e-12);
5064 }
5065
5066 #[test]
5067 fn pow_two_lowers_to_single_mul() {
5068 let e = pow(var(0), cnst(2.0));
5069 let t = Tape::build(&e);
5070 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5071 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5072 assert!((t.eval(&[3.0]) - 9.0).abs() < 1e-12);
5073 }
5074
5075 #[test]
5076 fn pow_three_lowers_to_two_muls() {
5077 let e = pow(var(0), cnst(3.0));
5078 let t = Tape::build(&e);
5079 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5080 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 2);
5081 assert!((t.eval(&[2.0]) - 8.0).abs() < 1e-12);
5082 }
5083
5084 #[test]
5085 fn pow_eight_lowers_to_three_muls() {
5086 let e = pow(var(0), cnst(8.0));
5088 let t = Tape::build(&e);
5089 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5090 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 3);
5091 assert!((t.eval(&[2.0]) - 256.0).abs() < 1e-12);
5092 }
5093
5094 #[test]
5095 fn pow_negative_two_lowers_to_div() {
5096 let e = pow(var(0), cnst(-2.0));
5098 let t = Tape::build(&e);
5099 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5100 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Div(..))), 1);
5101 assert!((t.eval(&[4.0]) - (1.0 / 16.0)).abs() < 1e-12);
5102 }
5103
5104 #[test]
5105 fn pow_large_const_stays_generic() {
5106 let e = pow(var(0), cnst(9.0));
5108 let t = Tape::build(&e);
5109 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5110 }
5111
5112 #[test]
5113 fn pow_non_integer_const_stays_generic() {
5114 let e = pow(var(0), cnst(1.5));
5116 let t = Tape::build(&e);
5117 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 1);
5118 }
5119
5120 #[test]
5121 fn pow_const_through_cse_const() {
5122 let two = Arc::new(cnst(2.0));
5124 let e = Expr::Binary(BinOp::Pow, Box::new(var(0)), Box::new(Expr::Cse(two)));
5125 let t = Tape::build(&e);
5126 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Pow(..))), 0);
5127 assert_eq!(count_op(&t, |o| matches!(o, TapeOp::Mul(..))), 1);
5128 }
5129
5130 #[test]
5131 fn hessian_pow_three_matches_fd() {
5132 let e = add(mul(cnst(5.0), pow(var(0), cnst(3.0))), mul(var(0), var(1)));
5134 let t = Tape::build(&e);
5135 fd_check(&t, &[1.7, 0.8], 2, 1e-5);
5136 }
5137
5138 #[test]
5139 fn hessian_pow_negative_matches_fd() {
5140 let e = add(pow(var(0), cnst(-2.0)), pow(var(1), cnst(2.0)));
5142 let t = Tape::build(&e);
5143 fd_check(&t, &[1.3, 2.4], 2, 1e-5);
5144 }
5145
5146 #[test]
5147 fn hessian_pow_half_matches_fd() {
5148 let e = add(pow(var(0), cnst(0.5)), mul(var(0), var(1)));
5150 let t = Tape::build(&e);
5151 fd_check(&t, &[2.5, 1.1], 2, 1e-5);
5152 }
5153
5154 #[test]
5155 fn hessian_sparsity_through_cse() {
5156 let body = Arc::new(add(var(0), var(1)));
5159 let e = add(
5160 pow(Expr::Cse(body.clone()), cnst(2.0)),
5161 Expr::Cse(body.clone()),
5162 );
5163 let t = Tape::build(&e);
5164 let s = t.hessian_sparsity();
5165 assert!(s.contains(&(0, 0)));
5166 assert!(s.contains(&(1, 0)));
5167 assert!(s.contains(&(1, 1)));
5168 assert_eq!(s.len(), 3);
5169 }
5170
5171 #[test]
5172 fn pow_forward_tangent_matches_reverse_gradient_at_base_zero() {
5173 let e = pow(var(0), var(1));
5182 let t = Tape::build(&e);
5183 assert!(
5186 t.ops.iter().any(|op| matches!(op, TapeOp::Pow(_, _))),
5187 "expected a Pow op in the tape; got {:?}",
5188 t.ops
5189 );
5190 let x = [0.0, 1.0];
5191 let n = t.ops.len();
5192
5193 let mut grad = vec![0.0; 2];
5195 t.gradient_seed(&x, 1.0, &mut grad);
5196
5197 let vals = t.forward(&x);
5199 let mut dot = vec![0.0; n];
5200 t.forward_tangent(&vals, 0, &mut dot);
5201 let fwd_dfx0 = dot[n - 1];
5202
5203 assert!(
5204 (grad[0] - 1.0).abs() < 1e-12,
5205 "reverse gradient df/dx0 at base 0 should be 1, got {}",
5206 grad[0]
5207 );
5208 assert!(
5209 (fwd_dfx0 - grad[0]).abs() < 1e-12,
5210 "forward tangent df/dx0 = {fwd_dfx0} must match reverse gradient {} at base 0",
5211 grad[0]
5212 );
5213 }
5214
5215 #[test]
5216 #[should_panic(expected = "external function calls are not supported on the")]
5217 fn hybrid_promoted_cse_with_funcall_reports_clear_message() {
5218 let body = Arc::new(Expr::Funcall {
5226 id: 0,
5227 args: vec![FuncallArg::Real(var(0))],
5228 });
5229 let exprs = vec![
5230 add(Expr::Cse(body.clone()), cnst(1.0)),
5231 add(Expr::Cse(body.clone()), cnst(2.0)),
5232 ];
5233 HybridTape::build_multi(&exprs);
5234 }
5235
5236 #[test]
5247 fn directional_hybrid_hessian_matches_flat_directional() {
5248 let body = Arc::new(unary(UnaryOp::Exp, mul(cnst(0.5), add(var(0), var(1)))));
5251 let exprs = vec![
5252 pow(Expr::Cse(body.clone()), cnst(2.0)),
5253 mul(Expr::Cse(body.clone()), var(2)),
5254 add(
5255 unary(UnaryOp::Sin, Expr::Cse(body.clone())),
5256 pow(var(2), cnst(2.0)),
5257 ),
5258 ];
5259 let weights = [1.25, -0.75, 2.5];
5260 let x = [0.3, -0.1, 0.7];
5261 let n = 3;
5262
5263 let hybrid = HybridTape::build_multi(&exprs);
5264 assert!(
5265 hybrid.n_prelude_ops() > 0,
5266 "a CSE shared by 3 roots must be promoted into the prelude"
5267 );
5268
5269 let seeds = [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]];
5271 for seed in &seeds {
5272 let mut flat_out = vec![0.0; n];
5273 for (e, &wt) in exprs.iter().zip(&weights) {
5274 let t = Tape::build(e);
5275 let vals = t.forward(&x);
5276 let m = t.ops.len();
5277 let (mut dot, mut adj, mut adj_dot) = (vec![0.0; m], vec![0.0; m], vec![0.0; m]);
5278 t.hessian_directional(
5279 &vals,
5280 seed,
5281 wt,
5282 &mut flat_out,
5283 &mut dot,
5284 &mut adj,
5285 &mut adj_dot,
5286 );
5287 }
5288
5289 let mut hyb_out = vec![0.0; n];
5290 let np = hybrid.n_prelude_ops();
5291 let ml = hybrid.max_summand_ops();
5292 let mut prelude_vals = vec![0.0; np];
5293 let mut prelude_dot = vec![0.0; np];
5294 let mut prelude_adj = vec![0.0; np];
5295 let mut prelude_adj_dot = vec![0.0; np];
5296 let (mut local_dot, mut local_adj, mut local_adj_dot) =
5297 (vec![0.0; ml], vec![0.0; ml], vec![0.0; ml]);
5298 let creach: Vec<u32> = {
5304 let mut u: BTreeSet<u32> = BTreeSet::new();
5305 for s in &hybrid.summands {
5306 u.extend(s.prelude_reach.iter().map(|&p| p as u32));
5307 }
5308 u.into_iter().collect()
5309 };
5310 hybrid.forward_prelude(&x, &mut prelude_vals);
5311 hybrid.prelude_tangent(&prelude_vals, seed, &creach, &mut prelude_dot);
5312 for (s, &wt) in hybrid.summands.iter().zip(&weights) {
5313 let mut local_vals = vec![0.0; s.ops.len()];
5314 hybrid.forward_summand(s, &x, &prelude_vals, &mut local_vals);
5315 hybrid.hessian_summand_directional(
5316 s,
5317 &local_vals,
5318 &prelude_dot,
5319 seed,
5320 wt,
5321 &mut hyb_out,
5322 &mut local_dot,
5323 &mut local_adj,
5324 &mut local_adj_dot,
5325 &mut prelude_adj,
5326 &mut prelude_adj_dot,
5327 );
5328 }
5329 hybrid.prelude_reverse_directional(
5330 &prelude_vals,
5331 &prelude_dot,
5332 &creach,
5333 &mut hyb_out,
5334 &mut prelude_adj,
5335 &mut prelude_adj_dot,
5336 );
5337
5338 for k in 0..n {
5339 let scale = flat_out[k].abs().max(1.0);
5340 assert!(
5341 (hyb_out[k] - flat_out[k]).abs() <= 1e-13 * scale,
5342 "seed {seed:?} entry {k}: hybrid {} vs flat {}",
5343 hyb_out[k],
5344 flat_out[k]
5345 );
5346 }
5347 assert!(
5348 hyb_out.iter().any(|v| *v != 0.0),
5349 "an all-zero H·s is no comparison"
5350 );
5351 assert!(prelude_adj.iter().all(|v| *v == 0.0));
5354 assert!(prelude_adj_dot.iter().all(|v| *v == 0.0));
5355 }
5356 }
5357}