Rename iter_outputs -> iter_embedded to avoid confusion

lukstafi · lukstafi · commit 7dd2d3591ce6 · 2024-10-05T22:40:08.000+02:00
diff --git a/bin/compilation_speed.ml b/bin/compilation_speed.ml
@@ -31,7 +31,7 @@ let benchmark_overhead backend () =
   let init_assign_x = link ctx @@ compile ~name:"init_assign_x" IDX.empty mock_update_x in
   let f_routine = link init_assign_x.context @@ compile IDX.empty update_f.fwd_bprop in
   Tensor.print_tree ~with_grad:true ~with_backend_info:true ~depth:9 f;
-  Tensor.iter_outputs f ~f:(fun a -> ignore (from_host f_routine.context a : bool));
+  Tensor.iter_embedded f ~f:(fun a -> ignore (from_host f_routine.context a : bool));
 
   let xs = Array.init n_data ~f:Float.(fun i -> of_int i - (of_int n_data /. 2.)) in
   let open Operation.At in
diff --git a/bin/zero2hero_1of7.ml b/bin/zero2hero_1of7.ml
@@ -163,9 +163,9 @@ let _suspended () =
   let device = new_virtual_device @@ get_device ~ordinal:0 in
   let update = Train.grad_update l in
   let routine = link (init device) @@ compile IDX.empty @@ update.fwd_bprop in
-  Tensor.iter_outputs l ~f:(fun a -> ignore (from_host routine.context a : bool));
+  Tensor.iter_embedded l ~f:(fun a -> ignore (from_host routine.context a : bool));
   Train.run routine;
-  Tensor.iter_outputs l ~f:(fun a -> ignore (to_host routine.context a : bool));
+  Tensor.iter_embedded l ~f:(fun a -> ignore (to_host routine.context a : bool));
   await device;
   Stdio.print_endline
     {|
@@ -177,7 +177,7 @@ let _suspended () =
     link routine.context @@ compile IDX.empty @@ Train.sgd_update ~learning_rate update
   in
   (* learning_rate is virtual so this will not print anything. *)
-  Tensor.iter_outputs learning_rate ~f:(fun a ->
+  Tensor.iter_embedded learning_rate ~f:(fun a ->
       ignore (from_host routine.context a : bool));
   Stdio.print_endline
     {|
@@ -187,7 +187,7 @@ let _suspended () =
   List.iter [ a.value; b.value; c.value; f.value ] ~f:(fun a ->
       assert (from_host routine.context a));
   Train.run routine;
-  Tensor.iter_outputs l ~f:(fun a -> ignore (to_host routine.context a : bool));
+  Tensor.iter_embedded l ~f:(fun a -> ignore (to_host routine.context a : bool));
   await device;
   Stdio.print_endline
     {|
@@ -198,7 +198,7 @@ let _suspended () =
   let update = Train.grad_update l in
   let routine = link routine.context @@ compile IDX.empty update.fwd_bprop in
   Train.run routine;
-  Tensor.iter_outputs l ~f:(fun a -> ignore (to_host routine.context a : bool));
+  Tensor.iter_embedded l ~f:(fun a -> ignore (to_host routine.context a : bool));
   await device;
   Stdio.print_endline
     {|
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -385,7 +385,7 @@ let rec inputs_and_outputs t =
   in
   (non_embedded - embedded, embedded)
 
-let iter_outputs ~f t = Set.iter ~f @@ snd @@ inputs_and_outputs t
+let iter_embedded ~f t = Set.iter ~f @@ snd @@ inputs_and_outputs t
 let input_nodes t = fst @@ inputs_and_outputs t
 let debug_name t = Tn.debug_name t.value
 let debug_grad t = Tn.debug_name (Option.value_exn t.diff).grad
diff --git a/lib/tensor.mli b/lib/tensor.mli
@@ -28,8 +28,8 @@ type t = {
           shape inference. *)
   children : subtensor list;
   non_embedded : tn_set;
-      (** These tensor nodes ([value], resp. [grad] of {!diff}) of the children which are not
-          computed by [forward], resp. [backprop] of {!diff}. *)
+      (** These tensor nodes ([value], resp. {!grad} of {!field-diff}) of the children which are not
+          computed by [forward], resp. {!backprop} of {!field-diff}. *)
 }
 [@@deriving sexp_of]
 (** Information needed for compositional code generation. *)
@@ -196,11 +196,11 @@ val consume_backprop_code : t -> asgns * asgns
 
 val input_nodes : t -> tn_set
 (** The nodes of descendant tensors whose computation is not embedded by the given tensor. They are
-    "inputs" coming from other computations. *)
+    "inputs" coming from other computations. NOTE: this a specific, narrow meaning of "inputs". *)
 
-val iter_outputs : f:(tn -> unit) -> t -> unit
-(** [iter_outputs t] iterates over all descendant nodes that are embedded, i.e. are not members
-    of [input_nodes t]. *)
+val iter_embedded : f:(tn -> unit) -> t -> unit
+(** [iter_embedded t] iterates over all descendant nodes that are embedded, i.e. are not members of
+    [input_nodes t] -- see {!input_nodes}. *)
 
 val unsafe_reinitialize : unit -> unit
 (** Bring global state to its initialization values. This invalidates any previously defined tensors
diff --git a/lib/train.ml b/lib/train.ml
@@ -262,18 +262,18 @@ let%track3_sexp round_robin_dry_run ~num_devices jitbs ~dry_sync : unit =
 let set_virtual (a : Tn.t) = Tn.update_memory_mode a Virtual 29
 
 let every_non_literal_on_host =
-  Tensor.iter_outputs ~f:(fun a ->
+  Tensor.iter_embedded ~f:(fun a ->
       if Tn.mode_is_unspecified a && not (Tn.known_constant a) then set_hosted a)
 
 let%debug2_sexp all_host_to_device (type context)
     (module Backend : Backend_type with type context = context) context =
   let f tn = ignore (Backend.from_host context tn : bool) in
-  Tensor.iter_outputs ~f
+  Tensor.iter_embedded ~f
 
 let%debug2_sexp all_device_to_host (type context)
     (module Backend : Backend_type with type context = context) context =
   let f tn = ignore (Backend.to_host context tn : bool) in
-  Tensor.iter_outputs ~f
+  Tensor.iter_embedded ~f
 
 (** Executes the jitted code and copies arrays embedded in the given tenosor from and to host,
     synchronizes before copying to host. If [looping] is provided, loops over bindings and executes
