update batch test

Author: ZuseZ4

tests/codegen/autodiffv.rs

@@ -1,13 +1,22 @@
 //@ compile-flags: -Zautodiff=Enable -C opt-level=3  -Clto=fat
 //@ no-prefer-dynamic
 //@ needs-enzyme
+
 #![feature(autodiff)]
 
 use std::autodiff::autodiff;
 
-#[autodiff(d_square, Reverse, 4, Duplicated, Active)]
+#[autodiff(d_square3, Forward, Dual, DualOnly)]
+#[no_mangle]
+fn squaref(x: &f32) -> f32 {
+    2.0 * x * x
+}
+
+
+#[autodiff(d_square2, Forward, 4, Dual, DualOnly)]
+#[autodiff(d_square, Forward, 4, Dual, Dual)]
 #[no_mangle]
-fn square(x: &f64) -> f64 {
+fn square(x: &f32) -> f32 {
     x * x
 }
 
@@ -33,21 +42,31 @@ fn square(x: &f64) -> f64 {
 // CHECK-NEXT:}
 
 fn main() {
-    let x = 3.0;
+    let x = std::hint::black_box(3.0);
     let output = square(&x);
+    dbg!(&output);
     assert_eq!(9.0, output);
+    dbg!(squaref(&x));
 
-    let mut df_dx1 = 0.0;
-    let mut df_dx2 = 0.0;
-    let mut df_dx3 = 0.0;
+    let mut df_dx1 = 1.0;
+    let mut df_dx2 = 2.0;
+    let mut df_dx3 = 3.0;
     let mut df_dx4 = 0.0;
-    let [o1, o2, o3, o4] = d_square(&x, &mut df_dx1, &mut df_dx2, &mut df_dx3, &mut df_dx4, 1.0);
-    assert_eq!(output, o1);
-    assert_eq!(output, o2);
-    assert_eq!(output, o3);
-    assert_eq!(output, o4);
-    assert_eq!(6.0, df_dx1);
-    assert_eq!(6.0, df_dx2);
-    assert_eq!(6.0, df_dx3);
-    assert_eq!(6.0, df_dx4);
+    let [o1,o2,o3,o4] = d_square2(&x, &mut df_dx1, &mut df_dx2, &mut df_dx3,  &mut df_dx4);
+    dbg!(o1, o2, o3, o4);
+    let [output2, o1,o2,o3,o4] = d_square(&x, &mut df_dx1, &mut df_dx2, &mut df_dx3,  &mut df_dx4);
+    dbg!(o1, o2, o3, o4);
+    assert_eq!(output, output2);
+    assert!((6.0 - o1).abs() < 1e-10);
+    assert!((12.0 - o2).abs() < 1e-10);
+    assert!((18.0 - o3).abs() < 1e-10);
+    assert!((0.0 - o4).abs() < 1e-10);
+    assert_eq!(1.0, df_dx1);
+    assert_eq!(2.0, df_dx2);
+    assert_eq!(3.0, df_dx3);
+    assert_eq!(0.0, df_dx4);
+    assert_eq!(d_square3(&x, &mut df_dx1), 2.0 * o1);
+    assert_eq!(d_square3(&x, &mut df_dx2), 2.0 * o2);
+    assert_eq!(d_square3(&x, &mut df_dx3), 2.0 * o3);
+    assert_eq!(d_square3(&x, &mut df_dx4), 2.0 * o4);
 }