Input/Output Operations

The Nx_io module provides functions to load and save Nx tensors in various file formats, including image formats and NumPy formats.

Features

Image Operations

Loading Images

Images can be loaded as uint8 tensors:

open Nx_io

(* Load as RGB: shape [|height; width; 3|] *)
let img = load_image "photo.png"

(* Load as grayscale: shape [|height; width|] *)
let gray = load_image ~grayscale:true "photo.png"

Saving Images

Save uint8 tensors as images:

(* Save RGB or grayscale based on shape *)
save_image img "output.png"

Supported shapes:

NumPy Format Support

Single Arrays (.npy)

The .npy format stores a single array with its dtype and shape information:

(* Load array with runtime-detected type *)
let P arr = load_npy "data.npy" in
(* arr : ('a, 'b) Nx.t *)

(* Convert to specific type *)
let float_arr = load_npy "data.npy" |> to_float32

(* Save array *)
save_npy my_array "output.npy"

Archives (.npz)

The .npz format stores multiple named arrays in a compressed archive:

(* Load entire archive *)
let archive = load_npz "bundle.npz" in

(* Access specific array *)
match Hashtbl.find_opt archive "weights" with
| Some (P arr) -> 
    let weights = to_float32 (P arr) in
    (* use weights *)
| None -> failwith "weights not found"

(* Load single array directly *)
let P data = load_npz_member ~path:"bundle.npz" ~name:"data"

(* Save multiple arrays *)
save_npz [
  ("inputs", P input_array);
  ("labels", P label_array);
  ("weights", P weight_array)
] "model.npz"

Packed Arrays and Type Conversions

Since file formats store type information that's only known at runtime, loaded arrays are wrapped in the Nx_io.packed_nx type:

type packed_nx = P : ('a, 'b) Nx.t -> packed_nx

Convert packed arrays to specific types using the provided functions:

let packed = load_npy "data.npy" in
let float32_array = to_float32 packed
let int32_array = to_int32 packed
let uint8_array = to_uint8 packed
(* etc. *)

Available conversions:

Examples

Image Processing Pipeline

open Nx
open Nx_io

let process_image input_path output_path =
  (* Load image *)
  let img = load_image input_path in
  
  (* Convert to float for processing *)
  let img_float = Nx.astype float32 img in
  
  (* Normalize to [0, 1] *)
  let normalized = Nx.div_s img_float 255.0 in
  
  (* Apply some processing *)
  let processed = 
    normalized
    |> Nx.mul_s 0.8  (* Reduce brightness *)
    |> Nx.add_s 0.1  (* Add bias *)
  in
  
  (* Convert back to uint8 *)
  let result = 
    processed
    |> Nx.mul_s 255.0
    |> Nx.clip ~min:0.0 ~max:255.0
    |> Nx.astype uint8
  in
  
  (* Save result *)
  save_image result output_path

Model Checkpoint Save/Load

let save_checkpoint ~path ~epoch ~model =
  let weights = Model.get_weights model in
  let optimizer_state = Model.get_optimizer_state model in
  
  save_npz [
    ("epoch", P (Nx.scalar int32 epoch));
    ("weights", P weights);
    ("optimizer_state", P optimizer_state);
  ] path

let load_checkpoint path =
  let archive = load_npz path in
  let epoch = 
    match Hashtbl.find_opt archive "epoch" with
    | Some p -> Nx.item [] (to_int32 p)
    | None -> failwith "epoch not found"
  in
  let weights = 
    match Hashtbl.find_opt archive "weights" with
    | Some p -> to_float32 p
    | None -> failwith "weights not found"
  in
  (epoch, weights)

Error Handling

All I/O operations may raise Failure exceptions:

Performance Considerations

See Also