Fix: Always reinitialize global state at the beginning of let%expect_test

lukstafi · lukstafi · commit f5fb81dd9a98 · 2024-09-16T22:45:44.000+02:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -29,6 +29,7 @@
 - Fixed loss of significant digits for small numbers when outputting files.
 - Added missing mixed-precision conversions in the `C_syntax` backend builder.
 - Restored the functionality of debug logging from the cuda backend.
+- Always reinitialize global state at the beginning of `let%expect_test`, to make them more deterministic.
 
 ## [0.4.0] -- 2024-09-04
 
diff --git a/lib/shape.ml b/lib/shape.ml
@@ -238,11 +238,11 @@ end
 
 type update_id = Update_id.t [@@deriving equal, compare, hash, sexp]
 
-let get_update_id =
-  let uid = ref 0 in
-  fun () ->
-    Int.incr uid;
-    Update_id.Update_id !uid
+let update_uid = ref 0
+
+let get_update_id () =
+  Int.incr update_uid;
+  Update_id.Update_id !update_uid
 
 type update_step = { shape : t; logic : logic; id : update_id } [@@deriving sexp]
 (** Data required for a shape inference update step. Ideally, an update should be performed at least
@@ -538,6 +538,12 @@ let state = ref Row.empty_env
 let active_update_steps = ref []
 let active_constraints = ref []
 
+let unsafe_reinitialize () =
+  update_uid := 0;
+  state := Row.empty_env;
+  active_update_steps := [];
+  active_constraints := []
+
 let iter_shapes update_step ~f =
   f update_step.shape;
   match update_step.logic with
diff --git a/lib/shape.mli b/lib/shape.mli
@@ -101,6 +101,9 @@ val to_string_hum :
   t ->
   string
 
+val unsafe_reinitialize : unit -> unit
+(** Bring global state to its initialization values. This invalidates any unfinished inference. *)
+
 (** {2 Internal-ish API.} *)
 
 (** How to propagate shape updates and do the last update of [Tensor.t.shape] when finalizing the
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -70,6 +70,13 @@ type session_state = {
 let session_state =
   { next_id = 0; forward_roots = Map.empty (module Int); backprop_roots = Map.empty (module Int) }
 
+let unsafe_reinitialize () =
+  session_state.next_id <- 0;
+  session_state.forward_roots <- Map.empty (module Int);
+  session_state.backprop_roots <- Map.empty (module Int);
+  Tn.Registry.clear Tn.registry;
+  Shape.unsafe_reinitialize ()
+
 let is_fwd_root t = Map.mem session_state.forward_roots t.id
 let remove_fwd_root t = session_state.forward_roots <- Map.remove session_state.forward_roots t.id
 let is_bprop_root t = Map.mem session_state.backprop_roots t.id
diff --git a/lib/tensor.mli b/lib/tensor.mli
@@ -188,6 +188,10 @@ val consume_backprop_code : t -> asgns * asgns
     [consume_backprop_code t] ensures [t] is a backprop root, removes it from backprop roots, and
     checks that there are no other backprop roots for tensors with children. *)
 
+val unsafe_reinitialize : unit -> unit
+(** Bring global state to its initialization values. This invalidates any previously defined tensors
+    and tensor nodes. Also invokes {!Shape.unsafe_reinitialize}. *)
+
 (** {2 Printing.} *)
 
 val header : t -> string
diff --git a/test/einsum_trivia.ml b/test/einsum_trivia.ml
diff --git a/test/hello_world_op.ml b/test/hello_world_op.ml
diff --git a/test/micrograd_demo.ml b/test/micrograd_demo.ml
@@ -7,6 +7,7 @@ module CDSL = Train.CDSL
 module Rand = Arrayjit.Rand.Lib
 
 let%expect_test "Micrograd README basic example" =
+  Tensor.unsafe_reinitialize ();
   Rand.init 0;
   let module Backend = (val Arrayjit.Backends.fresh_backend ()) in
   let backend = (module Backend : Train.Backend_type with type context = Backend.context) in
@@ -79,6 +80,7 @@ let%expect_test "Micrograd README basic example" =
     |}]
 
 let%expect_test "Micrograd half-moons example" =
+  Tensor.unsafe_reinitialize ();
   Rand.init 5;
   let module Backend = (val Arrayjit.Backends.fresh_backend ()) in
   let backend = (module Backend : Train.Backend_type with type context = Backend.context) in
diff --git a/test/moons_demo_parallel.ml b/test/moons_demo_parallel.ml
@@ -78,6 +78,7 @@ let main () =
   PrintBox_text.output Stdio.stdout plot_moons
 
 let%expect_test "Half-moons data parallel" =
+  Tensor.unsafe_reinitialize ();
   main ();
   (* NOTE: as of OCANNL 0.4, moons_demo_parallel, while deterministic on a single machine, gives
      slightly different results on machines with a different hardware, e.g. arm64, ppc. Here we list
diff --git a/test/zero2hero_1of7.ml b/test/zero2hero_1of7.ml
