Fix missing reshape in Tnode.create_with_reshape

lukstafi · lukstafi · commit dbff017fe1c3 · 2025-07-12T14:42:52.000+02:00
diff --git a/arrayjit/lib/ndarray.ml b/arrayjit/lib/ndarray.ml
@@ -385,15 +385,14 @@ let hash nd = Nativeint.hash (to_native nd)
 let hash_fold_t acc nd = hash_fold_nativeint acc (to_native nd)
 let hash_t nd = Nativeint.hash @@ to_native nd
 
-(** C function declarations for efficient copying *)
-external copy_with_padding_c : 
-  ('a, 'b) bigarray -> ('a, 'b) bigarray -> axis_padding array -> unit
+external copy_with_padding_c : ('a, 'b) bigarray -> ('a, 'b) bigarray -> axis_padding array -> unit
   = "arrayjit_copy_with_padding"
+(** C function declarations for efficient copying *)
 
+(** Copies the whole of [source] onto the parts of [target] skipping over padding margins --
+    requires that source dimensions + padding = target dimensions. *)
 let copy_with_padding ~source ~target ~padding =
-  let copy_impl source_arr target_arr =
-    copy_with_padding_c source_arr target_arr padding
-  in
+  let copy_impl source_arr target_arr = copy_with_padding_c source_arr target_arr padding in
   map2 { f2 = copy_impl } source target
 
 (** {2 *** Creating ***} *)
@@ -420,6 +419,11 @@ let%track7_sexp create_array ~debug:(_debug : string) (prec : Ops.prec) ~(dims :
       [%log _debug, ptr_to_string_hum result]]];
   result
 
+(** See {!Bigarray.reshape}. *)
+let reshape nd dims =
+  let f prec arr = as_array prec @@ Bigarray.reshape arr dims in
+  map_with_prec { f } nd
+
 let get_used_memory () = Atomic.get used_memory
 
 (** {2 *** Printing ***} *)
diff --git a/arrayjit/lib/tnode.ml b/arrayjit/lib/tnode.ml
@@ -634,35 +634,31 @@ let create_with_reshape ~id ~label ~base_ndarray ~dims ~padding ~from_padded ()
              Nd.as_array prec (Bigarray.reshape arr target_dims)
            in
            Some (Nd.map_with_prec { f = f_reshape_with_prec } base_ndarray)
-       | Some (padding_spec, _), false ->
+       | Some (padding, _), false ->
            (* Create new bigarray with padding and copy source into non-padding parts *)
-           let target_array =
-             Nd.create_array ~debug prec_val ~dims:target_dims ~padding:target_padding
-           in
+           let target = Nd.create_array ~debug prec_val ~dims:target_dims ~padding:target_padding in
            let source_dims = Nd.dims base_ndarray in
-           let copy_with_padding () =
-             (* Calculate actual data dimensions (target dims minus padding) *)
-             let data_dims =
-               Array.map2_exn target_dims padding_spec ~f:(fun dim { Nd.left; right } ->
-                   dim - left - right)
-             in
-             (* Check total elements match, allowing shape differences *)
-             let source_total = 
-               if Array.is_empty source_dims then 0 else Array.reduce_exn source_dims ~f:( * ) 
-             in
-             let data_total = 
-               if Array.is_empty data_dims then 0 else Array.reduce_exn data_dims ~f:( * ) 
-             in
-             if source_total <> data_total then
-               invalid_arg
-                 [%string
-                   "create_with_reshape: source has %{source_total#Int} elements but target data \
-                    area has %{data_total#Int} elements"];
-             (* Use C function for efficient copying *)
-             Nd.copy_with_padding ~source:base_ndarray ~target:target_array ~padding:padding_spec
+           (* Calculate actual data dimensions (target dims minus padding) *)
+           let data_dims =
+             Array.map2_exn target_dims padding ~f:(fun dim { Nd.left; right } ->
+                 dim - left - right)
+           in
+           (* Check total elements match, allowing shape differences *)
+           let source_total =
+             if Array.is_empty source_dims then 0 else Array.reduce_exn source_dims ~f:( * )
+           in
+           let data_total =
+             if Array.is_empty data_dims then 0 else Array.reduce_exn data_dims ~f:( * )
            in
-           copy_with_padding ();
-           Some target_array)
+           if source_total <> data_total then
+             invalid_arg
+               [%string
+                 "create_with_reshape: source has %{source_total#Int} elements but target data \
+                  area has %{data_total#Int} elements"];
+           (* Use C function for efficient copying *)
+           let source = Nd.reshape base_ndarray data_dims in
+           Nd.copy_with_padding ~source ~target ~padding;
+           Some target)
   and prec = lazy prec_val
   and size_in_bytes = lazy (num_elems tn * Ops.prec_in_bytes (Lazy.force tn.prec))
   and tn =
