nx_io: add save_txt/load_txt for 1D/2D tensors

CHANGES.md

@@ -9,6 +9,7 @@ All notable changes to this project will be documented in this file.
 
 ### Nx
 
+- Add `Nx_io.save_txt` / `Nx_io.load_txt` with NumPy-compatible formatting, comments, and dtype support (#120, @six-shot)
 - Clarify `reshape` documentation to match its view-only semantics (@tmattio)
 - Provide `nx.top`, `rune.top`, and `hugin.top` libraries that auto-install pretty printers in the OCaml toplevel and update Quill to load them (@tmattio)
 - Add `ifill` for explicit in-place fills and make `fill` return a copied tensor (@tmattio)

nx/lib/io/nx_io.ml

@@ -100,6 +100,10 @@ module Safe = struct
   let save_npz ?(overwrite = true) path items =
     Nx_npy.save_npz ~overwrite path items
 
+  (* Text I/O *)
+  let save_txt = Nx_txt.save
+  let load_txt = Nx_txt.load
+
   (* Conversions from packed arrays *)
 
   let as_float16 = Packed_nx.as_float16
@@ -164,3 +168,10 @@ let load_safetensor path = Safe.load_safetensor path |> unwrap_result
 
 let save_safetensor ?overwrite path items =
   Safe.save_safetensor ?overwrite path items |> unwrap_result
+
+let save_txt ?sep ?append ?newline ?header ?footer ?comments ~out arr =
+  Safe.save_txt ?sep ?append ?newline ?header ?footer ?comments ~out arr
+  |> unwrap_result
+
+let load_txt ?sep ?comments ?skiprows ?max_rows dtype path =
+  Safe.load_txt ?sep ?comments ?skiprows ?max_rows dtype path |> unwrap_result

nx/lib/io/nx_io.mli

@@ -191,6 +191,44 @@ val save_npz : ?overwrite:bool -> string -> (string * packed_nx) list -> unit
 
     Save multiple named nxs to a NumPy `.npz` archive. *)
 
+(** {1 Text Format} *)
+
+val save_txt :
+  ?sep:string ->
+  ?append:bool ->
+  ?newline:string ->
+  ?header:string ->
+  ?footer:string ->
+  ?comments:string ->
+  out:string ->
+  ('a, 'b) Nx.t ->
+  unit
+(** [save_txt ?sep ?append ?newline ?header ?footer ?comments ~out t]
+
+    Save a scalar, 1D, or 2D tensor to a text file. Each row is written on a new
+    line, values separated by [sep] (default: a single space). If [append] is
+    [true], data is appended to [out]; otherwise the file is truncated/created.
+    Optional [header] and [footer] strings are emitted before and after the
+    data, prefixed with [comments] (default: ["# "]). The [newline] argument
+    controls the end-of-line separator. Only numeric and boolean dtypes are
+    supported. *)
+
+val load_txt :
+  ?sep:string ->
+  ?comments:string ->
+  ?skiprows:int ->
+  ?max_rows:int ->
+  ('a, 'b) Nx.dtype ->
+  string ->
+  ('a, 'b) Nx.t
+(** [load_txt ?sep ?comments ?skiprows ?max_rows dtype path]
+
+    Load a tensor from a text file previously written by [save_txt] or a
+    compatible format. Lines beginning with [comments] (after trimming leading
+    whitespace) are ignored. Leading [skiprows] raw lines are skipped. When data
+    has a single row or a single column, the result is 1D; otherwise a 2D tensor
+    of shape [|rows; cols|] is returned. *)
+
 (** {1 SafeTensors Format} *)
 
 val load_safetensor : string -> archive
@@ -288,6 +326,30 @@ module Safe : sig
     (unit, error) result
   (** Safe alias for [save_npz] *)
 
+  (** {2 Text Format} *)
+
+  val save_txt :
+    ?sep:string ->
+    ?append:bool ->
+    ?newline:string ->
+    ?header:string ->
+    ?footer:string ->
+    ?comments:string ->
+    out:string ->
+    ('a, 'b) Nx.t ->
+    (unit, error) result
+  (** Safe alias for [save_txt] *)
+
+  val load_txt :
+    ?sep:string ->
+    ?comments:string ->
+    ?skiprows:int ->
+    ?max_rows:int ->
+    ('a, 'b) Nx.dtype ->
+    string ->
+    (('a, 'b) Nx.t, error) result
+  (** Safe alias for [load_txt] *)
+
   (** {2 SafeTensors Format} *)
 
   val load_safetensor : string -> (archive, error) result