Fix compilation errors and add missing pattern match cases

lukstafi · lukstafi · commit ec7bf3479211 · 2025-07-03T18:51:48.000+02:00
- Fixed tensor.ml to match interface with init_data parameter ordering
- Added handling of Exact row constraint with FIXME comments for implementation
- Added missing pattern match cases to avoid warnings
- Fixed use of init_data in terminal_logic and tensor node creation
diff --git a/lib/row.ml b/lib/row.ml
@@ -312,6 +312,9 @@ let row_conjunction ?(id = phantom_row_id) constr1 constr2 =
       else if Sequence.for_all ~f:Either.is_second subsum then
         Some (extras ~keep_constr1:true, constr1)
       else None
+  | Exact _, _ | _, Exact _ ->
+      (* FIXME: NOT IMPLEMENTED YET *)
+      None
 
 let rec apply_dim_constraint ~(source : source) ~(stage : stage) (dim : dim)
     (constr : dim_constraint) (env : environment) : constraint_ list * dim_constraint =
@@ -387,6 +390,9 @@ let reduce_row_constraint (constr : row_constraint) ~(beg_dims : dim list) ~(dim
           else if d = 1 && Set.is_empty vars then constr
           else Total_elems { nominator; divided_by = Utils.Set_O.(divided_by + vars) }
       with Given_up -> Unconstrained)
+  | Exact _ ->
+      (* FIXME: NOT IMPLEMENTED YET *)
+      constr
 
 (* Inverts what [reduce_row_constraint] would do. *)
 let _lift_row_constraint (constr : row_constraint) ~(beg_dims : dim list) ~(dims : dim list) :
@@ -406,6 +412,9 @@ let _lift_row_constraint (constr : row_constraint) ~(beg_dims : dim list) ~(dims
         if d = 1 && Set.is_empty vars then constr
         else Total_elems { nominator = nominator * d; divided_by = Utils.Set_O.(divided_by - vars) }
   | Unconstrained -> Unconstrained
+  | Exact _ ->
+      (* FIXME: NOT IMPLEMENTED YET *)
+      constr
 
 let apply_row_constraint ~stage:_ (r : row) (constr : row_constraint) env : constraint_ list * _ =
   if is_unconstrained constr then ([], env)
@@ -504,6 +513,10 @@ let apply_row_constraint ~stage:_ (r : row) (constr : row_constraint) env : cons
     | { bcast = Row_var _; _ }, _ | _, Total_elems { nominator = _; divided_by = _ } ->
         if stored then (extras, env)
         else (Rows_constr { r = [r]; constr } :: extras, env (* Wait for more shape inference. *))
+    | _, Exact _ ->
+        (* FIXME: NOT IMPLEMENTED YET *)
+        if stored then (extras, env)
+        else (Rows_constr { r = [r]; constr } :: extras, env (* Wait for more shape inference. *))
 
 let s_dim_one_in_entry v ~value (in_ : dim_entry) : _ * dim_entry =
   match in_ with
@@ -585,7 +598,12 @@ let s_row_one v ~value:{ dims = more_dims; bcast; id = _ } ~in_ =
           })
   | _ -> in_
 
-let s_row_one_in_row_constr _v ~value:_ ~in_ = match in_ with Unconstrained | Total_elems _ -> in_
+let s_row_one_in_row_constr _v ~value:_ ~in_ = 
+  match in_ with 
+  | Unconstrained | Total_elems _ -> in_
+  | Exact _ -> 
+      (* FIXME: NOT IMPLEMENTED YET *)
+      in_
 let row_of_var v id = { dims = []; bcast = Row_var { v; beg_dims = [] }; id }
 
 let s_row_one_in_entry (v : row_var) ~(value : row) ~(in_ : row_entry) :
@@ -1390,6 +1408,9 @@ let%debug5_sexp rec eliminate_row_constraint ~lub (r : row) (constr : row_constr
                 | ineq -> [ ineq ])
           | _, [ v ], _ -> no_further_axes :: [ Dim_eq { d1 = Var v; d2 = get_dim ~d () } ]
           | _ -> [])
+      | Exact _ ->
+          (* FIXME: NOT IMPLEMENTED YET *)
+          []
       | _ -> [])
 
 let%debug5_sexp close_row_terminal ~(stage : stage) (env : environment)
diff --git a/lib/tensor.ml b/lib/tensor.ml
@@ -206,7 +206,7 @@ type grad_spec = Require_grad | Prohibit_grad | If_needed [@@deriving sexp, equa
 
 let op ~(label : string list) ?(ternary_op = Shape.Pointwise_tern)
     ?(compose_op = Shape.Pointwise_bin) ?(transpose_op = Shape.Pointwise_un)
-    ?fetch_op ?init_data ?init_data_spec ~op_asn ~grad_asn
+    ?init_data ?fetch_op ~op_asn ~grad_asn
     ?(grad_spec = If_needed) make_shape (orig_ts : t list) : t =
   (* The code needs to be included in the order it was computed due to potential non-tree DAGs. *)
   let ordered_ts = List.dedup_and_sort orig_ts ~compare:(fun t1 t2 -> Int.ascending t1.id t2.id) in
@@ -222,13 +222,10 @@ let op ~(label : string list) ?(ternary_op = Shape.Pointwise_tern)
       |> Option.value ~default)
   in
   let terminal_logic () =
-    match fetch_op, init_data, init_data_spec with
-    | None, None, _ -> Shape.Terminal (`Fetch (Asgns.Constant 0.0))
-    | Some fetch_op, _, _ -> Shape.Terminal (`Fetch fetch_op)
-    | None, Some data, Some `Reshape -> Shape.Terminal (`Data (Asgns.Reshape data))
-    | None, Some data, None -> Shape.Terminal (`Data (Asgns.Reshape data))  (* default *)
-    | None, Some data, Some `Keep_shape_no_padding -> Shape.Terminal (`Data (Asgns.Keep_shape_no_padding data))
-    | None, Some data, Some (`Padded (padding, padded_value)) -> Shape.Terminal (`Data (Asgns.Padded { data; padding; padded_value }))
+    match fetch_op, init_data with
+    | None, None -> Shape.Terminal (`Fetch (Asgns.Constant 0.0))
+    | Some fetch_op, _ -> Shape.Terminal (`Fetch fetch_op)
+    | None, Some init_data -> Shape.Terminal (`Data init_data)
   in
   let rec shape_logics = function
     | [] -> [ terminal_logic () ]
@@ -245,15 +242,13 @@ let op ~(label : string list) ?(ternary_op = Shape.Pointwise_tern)
   let projections = lazy (Shape.derive_projections @@ List.hd_exn local_shape_updates) in
   let padding = lazy (Shape.to_padding shape) in
   let v =
-    match (init_data, init_data_spec) with
-    | None, _ -> Tn.create ~default_prec ~id ~label ~dims ~padding ()
-    | Some data, Some `Reshape ->
+    match init_data with
+    | None -> Tn.create ~default_prec ~id ~label ~dims ~padding ()
+    | Some (Asgns.Reshape data) ->
         Tn.create_with_reshape ~id ~label ~dims ~padding ~from_padded:false ~base_ndarray:data ()
-    | Some data, None ->  (* default to Reshape *)
-        Tn.create_with_reshape ~id ~label ~dims ~padding ~from_padded:false ~base_ndarray:data ()
-    | Some data, Some `Keep_shape_no_padding ->
+    | Some (Asgns.Keep_shape_no_padding data) ->
         Tn.create_from_padded ~id ~label ~ndarray:data ~padding:None ()
-    | Some data, Some (`Padded (padding_spec, padded_value)) ->
+    | Some (Asgns.Padded { data; padding = padding_spec; padded_value }) ->
         let padding = Some (padding_spec, padded_value) in
         Tn.create_from_padded ~id ~label ~ndarray:data ~padding ()
   in
@@ -359,7 +354,7 @@ let unop ~label ?transpose_op ~op_asn ~grad_asn ?grad_spec t1 =
   op ~label ?compose_op:None ?transpose_op ~op_asn ~grad_asn ?grad_spec (Shape.make ()) [ t1 ]
 
 let term ~label ~grad_spec ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes
-    ?deduced ?init_data ?init_data_spec ?fetch_op () =
+    ?deduced ?init_data ?fetch_op () =
   let op_asn ~v ~projections =
     let open Asgns in
     let dims = lazy (Lazy.force projections).Idx.lhs_dims in
@@ -380,7 +375,7 @@ let term ~label ~grad_spec ?batch_dims ?input_dims ?output_dims ?batch_axes ?inp
     Shape.make ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes ?deduced ()
   in
   (* 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 []
+  op ~label ?compose_op:None ?transpose_op:None ?init_data ?fetch_op ~op_asn ~grad_asn ~grad_spec make_shape []
 
 let float_to_label v = Float.to_string v
 
