Missing param initializations: zero2hero_1of7 test

lukstafi · lukstafi · commit e22b6f93bad5 · 2025-07-13T12:23:36.000+02:00
diff --git a/test/operations/zero2hero_1of7.ml b/test/operations/zero2hero_1of7.ml
@@ -51,8 +51,10 @@ let%expect_test "Graph drawing recompile" =
     |}];
   let%op f = (3 *. ("x" [ 5 ] **. 2)) - (4 *. x) + 5 in
   Train.every_non_literal_on_host f;
+  let f_init = Train.to_routine (module Backend) ctx IDX.empty @@ Tensor.init_params f in
   let f_upd = Train.grad_update f in
-  let f_bprop = Train.to_routine (module Backend) ctx IDX.empty f_upd in
+  let f_bprop = Train.to_routine (module Backend) f_init.context IDX.empty f_upd in
+  Train.run f_init;
   Train.run f_bprop;
   Tensor.print_tree ~with_grad:true ~depth:9 f;
   [%expect
@@ -267,9 +269,10 @@ let%expect_test "Simple gradients hosted" =
   Train.every_non_literal_on_host l;
   Train.every_non_literal_on_host learning_rate;
   let sgd = Train.sgd_update ~learning_rate l in
-  let grad_routine = Train.to_routine (module Backend) ctx IDX.empty grad in
+  let f_init = Train.to_routine (module Backend) ctx IDX.empty @@ Tensor.init_params l in
+  let grad_routine = Train.to_routine (module Backend) f_init.context IDX.empty grad in
   let sgd_routine = Train.to_routine (module Backend) grad_routine.context IDX.empty sgd in
-  (* Check out the initial state without running a forward pass. *)
+  (* Check out the initial state without running an init or forward pass. *)
   Tensor.print_tree ~spy:true ~with_grad:true ~depth:9 l;
   [%expect
     {|
@@ -292,6 +295,7 @@ let%expect_test "Simple gradients hosted" =
     |}];
   (* Do not update the params: all values and gradients will be at initial points, which are
      specified in the tensor in the brackets. *)
+  Train.run f_init;
   Train.run grad_routine;
   Tensor.print_tree ~with_grad:true ~depth:9 l;
   [%expect
@@ -398,7 +402,8 @@ let%expect_test "Simple gradients virtual" =
     #1 grad_a Material/28 │#3 grad_b Material/28 │                      │
     <not-in-yet>          │<not-in-yet>          │                      │
     |}];
-  let grad_routine = Train.to_routine (module Backend) ctx IDX.empty grad in
+  let f_init = Train.to_routine (module Backend) ctx IDX.empty @@ Tensor.init_params l in
+  let grad_routine = Train.to_routine (module Backend) f_init.context IDX.empty grad in
   (* Check out the state without running a forward pass or compiling the SGD update. *)
   Tensor.print_tree ~spy:true ~with_grad:true ~depth:9 l;
   [%expect
@@ -422,6 +427,7 @@ let%expect_test "Simple gradients virtual" =
     |}];
   (* Do not update the params: all values and gradients will be at initial points, which are
      specified in the tensor in the brackets. *)
+  Train.run f_init;
   Train.run grad_routine;
   Tensor.print_tree ~with_grad:true ~depth:9 l;
   [%expect
@@ -507,7 +513,9 @@ let%expect_test "2D neuron hosted" =
   let%op v = ("w" [ (-3, 1) ] * "x" [ 2; 0 ]) + "b" [ 6.7 ] in
   Train.every_non_literal_on_host v;
   let update = Train.grad_update v in
-  let routine = Train.to_routine (module Backend) ctx IDX.empty update in
+  let f_init = Train.to_routine (module Backend) ctx IDX.empty @@ Tensor.init_params v in
+  let routine = Train.to_routine (module Backend) f_init.context IDX.empty update in
+  Train.run f_init;
   Train.run routine;
   Tensor.print_tree ~with_grad:true ~depth:9 v;
   [%expect
@@ -534,7 +542,9 @@ let%expect_test "2D neuron virtual" =
   let ctx = Backend.make_context stream in
   let%op v = ("w" [ (-3, 1) ] * "x" [ 2; 0 ]) + "b" [ 6.7 ] in
   let update = Train.grad_update v in
-  let routine = Train.to_routine (module Backend) ctx IDX.empty update in
+  let f_init = Train.to_routine (module Backend) ctx IDX.empty @@ Tensor.init_params v in
+  let routine = Train.to_routine (module Backend) f_init.context IDX.empty update in
+  Train.run f_init;
   Train.run routine;
   Tensor.print_tree ~with_grad:true ~depth:9 v;
   [%expect
