Cleanup: remove unused ~v input to fetch_op

lukstafi · lukstafi · commit f6776d1687e3 · 2025-07-02T22:06:10.000+02:00
diff --git a/bin/primitive_ops.ml b/bin/primitive_ops.ml
@@ -27,7 +27,7 @@ let%debug_sexp graph_t () : unit =
   let xs = Array.init size ~f:Float.(fun i -> (of_int i / 10.) + 0.1) in
   let x_flat =
     Tensor.term ~grad_spec:Require_grad ~label:[ "x_flat" ]
-      ~fetch_op:(fun ~v:_ -> Constant_fill xs)
+      ~fetch_op:(Constant_fill xs)
       ()
   in
   let step_sym, bindings = IDX.get_static_symbol ~static_range:size IDX.empty in
diff --git a/bin/zero2hero_1of7.ml b/bin/zero2hero_1of7.ml
@@ -57,7 +57,7 @@ let _suspended () =
   let x_flat =
     Tensor.term ~grad_spec:Tensor.Require_grad
       ~label:[ "x_flat" ] (* ~input_dims:[] ~output_dims:[ 1 ] *)
-      ~fetch_op:(fun ~v:_ -> Constant_fill values)
+      ~fetch_op:(Constant_fill values)
       ()
   in
   let step_sym, bindings = IDX.get_static_symbol ~static_range:size IDX.empty in
@@ -111,7 +111,7 @@ let _suspended () =
   (* Yay, the whole shape gets inferred! *)
   let x_flat =
     Tensor.term ~grad_spec:Require_grad ~label:[ "x_flat" ]
-      ~fetch_op:(fun ~v:_ -> Constant_fill xs)
+      ~fetch_op:(Constant_fill xs)
       ()
   in
   let step_sym, bindings = IDX.get_static_symbol ~static_range:size IDX.empty in
diff --git a/lib/operation.ml b/lib/operation.ml
@@ -319,8 +319,7 @@ let range ?(label = []) ?(grad_spec = Tensor.Prohibit_grad) ?axis_label upto =
   let result =
     Tensor.term
       ~label:(("0" ^ "..." ^ Int.to_string upto) :: label)
-      ~grad_spec ~batch_dims:[] ~input_dims:[]
-      ~fetch_op:(fun ~v:_ -> Range_over_offsets)
+      ~grad_spec ~batch_dims:[] ~input_dims:[] ~fetch_op:Range_over_offsets
   in
   match axis_label with
   | None -> result ~output_dims:[ upto + 1 ] ()
@@ -344,8 +343,7 @@ let range_of_shape ?(label = []) ?(grad_spec = Tensor.Prohibit_grad) ?batch_dims
   Tensor.term
     ~label:(("r" ^ Idx.dims_to_string dims) :: label)
     ~grad_spec ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes
-    ~fetch_op:(fun ~v:_ -> Range_over_offsets)
-    ()
+    ~fetch_op:Range_over_offsets ()
 
 (** A [stop_gradient] is an identity in the forward pass and a no-op in the backprop pass. *)
 let stop_gradient ?(label = []) =
@@ -390,7 +388,7 @@ let random_seed =
   let seed = Option.value ~default:42 @@ Utils.settings.fixed_state_for_init in
   let res =
     Tensor.term ~label:[ "random_seed" ] ~grad_spec:Prohibit_grad
-      ~fetch_op:(fun ~v:_ -> Asgns.Constant_fill [| Int.to_float seed |])
+      ~fetch_op:(Asgns.Constant_fill [| Int.to_float seed |])
       ()
   in
   Tn.update_memory_mode res.value Tn.Effectively_constant 24;
@@ -469,14 +467,12 @@ module TDSL = struct
   *)
   let init_const ~l ?b ?(i = []) ~o values =
     Tensor.term ~label:[ l ] ~grad_spec:Prohibit_grad ?batch_dims:b ~input_dims:i ~output_dims:o
-      ~fetch_op:(fun ~v:_ -> Asgns.Constant_fill values)
-      ()
+      ~fetch_op:(Asgns.Constant_fill values) ()
 
   (** It's like `Tensor.param` but without shape inference. *)
   let init_param ~l ?(b = []) ?(i = []) ?(o = []) values =
     Tensor.term ~label:[ l ] ~grad_spec:Require_grad ~batch_dims:b ~input_dims:i ~output_dims:o
-      ~fetch_op:(fun ~v:_ -> Asgns.Constant_fill values)
-      ()
+      ~fetch_op:(Asgns.Constant_fill values) ()
 end
 
 module NTDSL = struct
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -343,32 +343,29 @@ let term ~label ~grad_spec ?batch_dims ?input_dims ?output_dims ?batch_axes ?inp
     let dims = lazy (Lazy.force projections).Idx.lhs_dims in
     match fetch_op with
     | None -> Asgns.empty_comp
-    | Some fetch_op_fn ->
-        let fetch_op = fetch_op_fn ~v in
-        (match fetch_op with
-        | Constant _ | Slice _ | Embed_symbol _ | Range_over_offsets | Constant_fill _
-        | Access (Uint4x32_to_prec_uniform _) ->
-            (* For these operations it makes sense to have a local / virtual tensor if the result is
-               consumed in the same computation. *)
-            ()
-        | Access _ ->
-            (* Note: [Access] can be used for merging across devices. But, some use cases of
-               [Access] will require a hosted tensor node. *)
-            Tn.update_memory_mode v Materialized 22);
+    | Some
+        (( Constant _ | Slice _ | Embed_symbol _ | Range_over_offsets | Constant_fill _
+         | Access (Uint4x32_to_prec_uniform _) ) as fetch_op) ->
+        Asgns.to_comp @@ Fetch { array = v; fetch_op; dims }
+    | Some (Access _ as fetch_op) ->
+        (* Note: [Access] can be used for merging across devices. But, some use cases of [Access]
+           will require a hosted tensor node. *)
+        Tn.update_memory_mode v Materialized 22;
         Asgns.to_comp @@ Fetch { array = v; fetch_op; dims }
   in
   let grad_asn ~t:_ ~g:_ ~projections:_ = Asgns.empty_comp in
   let make_shape =
     Shape.make ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes ?deduced ()
   in
-  op ~label ?compose_op:None ?transpose_op:None ~op_asn ~grad_asn ~grad_spec make_shape []
+  (* Note: fetch_op in op is used only for shape inference. *)
+  op ~label ?compose_op:None ?transpose_op:None ?fetch_op ~op_asn ~grad_asn ~grad_spec make_shape []
 
 let float_to_label v = Float.to_string v
 
 let number ?(label = []) ?axis_label ?(grad_spec = Prohibit_grad) c =
   (* Note: no axis label so that we do not conflict with user labels. *)
   let label = float_to_label c :: label in
-  let fetch_op ~v:_ = Ir.Assignments.Constant c in
+  let fetch_op = Ir.Assignments.Constant c in
   let t = term ~label ~grad_spec ~batch_dims:[] ~input_dims:[] ~fetch_op in
   let t =
     match axis_label with
@@ -416,7 +413,7 @@ let ndarray ?(label = []) ?(grad_spec = Prohibit_grad) ?batch_dims ?input_dims ?
   let t =
     term ~label ~grad_spec ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes
       ~deduced:Not_constrained
-      ~fetch_op:(fun ~v:_ -> Asgns.Constant_fill values)
+      ~fetch_op:(Asgns.Constant_fill values)
       ()
   in
   Tn.update_memory_mode t.value Effectively_constant 24;
@@ -428,7 +425,7 @@ let ndarray ?(label = []) ?(grad_spec = Prohibit_grad) ?batch_dims ?input_dims ?
 
 let param ?(more_label = []) ?input_dims ?output_dims ?input_axes ?output_axes ?deduced ?value
     ?values label =
-  let fetch_op_fn ~v:_ =
+  let fetch_op =
     match (values, value) with
     | Some values, None -> Asgns.Constant_fill values
     | None, Some value -> Asgns.Constant value
@@ -437,7 +434,7 @@ let param ?(more_label = []) ?input_dims ?output_dims ?input_axes ?output_axes ?
   in
   let t =
     term ~label:(label :: more_label) ~grad_spec:Require_grad ~batch_dims:[] ?input_dims
-      ?output_dims ?input_axes ?output_axes ?deduced ~fetch_op:fetch_op_fn ()
+      ?output_dims ?input_axes ?output_axes ?deduced ~fetch_op ()
   in
   let v = t.value in
   (* It is convenient to use the param syntax for volatiles (mutable embedded_nodes). *)
diff --git a/lib/tensor.mli b/lib/tensor.mli
@@ -181,7 +181,7 @@ val term :
   ?input_axes:(string * int) list ->
   ?output_axes:(string * int) list ->
   ?deduced:Shape.deduce_within_shape ->
-  ?fetch_op:(v:tn -> fetch_op) ->
+  ?fetch_op:fetch_op ->
   unit ->
   t
 (** A terminal: a constant, a parameter, an input of the model. The semantics of shape specification
diff --git a/test/einsum/moons_demo_variant.ml b/test/einsum/moons_demo_variant.ml
@@ -25,7 +25,6 @@ let () =
   let steps = epochs * 2 * len / batch_size in
   let noise () = Rand.float_range (-0.1) 0.1 in
   let moons_flat =
-    Bigarray.Genarray.init in
     Array.concat_map (Array.create ~len ())
       ~f:
         Float.(
diff --git a/test/primitive_ops.ml b/test/primitive_ops.ml
@@ -23,7 +23,7 @@ let plot_unop ~f ?(x_min = -5.) ?(x_max = 5.) () =
   in
   let x_flat =
     Tensor.term ~grad_spec:Require_grad ~label:[ "x_flat" ]
-      ~fetch_op:(fun ~v:_ -> Constant_fill xs)
+      ~fetch_op:(Constant_fill xs)
       ()
   in
   let step_sym, bindings = IDX.get_static_symbol ~static_range:size IDX.empty in
diff --git a/test/zero2hero_1of7.ml b/test/zero2hero_1of7.ml
@@ -174,7 +174,7 @@ let%expect_test "Graph drawing fetch" =
   (* Yay, the whole shape gets inferred! *)
   let x_flat =
     Tensor.term ~grad_spec:Require_grad ~label:[ "x_flat" ]
-      ~fetch_op:(fun ~v:_ -> Constant_fill xs)
+      ~fetch_op:(Constant_fill xs)
       ()
   in
   let step_sym, bindings = IDX.get_static_symbol ~static_range:size IDX.empty in

Original file line number	Diff line number	Diff line change
`@@ -181,7 +181,7 @@ val term :`
`181`	`181`	`?input_axes:(string * int) list ->`
`182`	`182`	`?output_axes:(string * int) list ->`
`183`	`183`	`?deduced:Shape.deduce_within_shape ->`
`184`		`- ?fetch_op:(v:tn -> fetch_op) ->`
	`184`	`+ ?fetch_op:fetch_op ->`
`185`	`185`	`unit ->`
`186`	`186`	`t`
`187`	`187`	`(** A terminal: a constant, a parameter, an input of the model. The semantics of shape specification`