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 I/O: Load and save images in PNG, JPEG, BMP, TGA, and GIF formats
• NumPy .npy format: Load and save single arrays in NumPy's native format
• NumPy .npz archives: Load and save multiple named arrays in compressed archives
• Runtime dtype detection: Handle arrays with types determined at runtime
• Type conversion utilities: Convert between different numeric types

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 "output.png" img

Supported shapes:

• [|height; width|] — Grayscale
• [|height; width; 1|] — Grayscale with explicit channel
• [|height; width; 3|] — RGB
• [|height; width; 4|] — RGBA

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"
(* arr : ('a, 'b) Nx.t *)

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

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

Archives (.npz)

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

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

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

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

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

Packed Arrays and Type Conversions

Since file formats store type information that's only known at runtime, loaded arrays are wrapped in the 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"
let float32_array = as_float32 packed
let int32_array = as_int32 packed
let uint8_array = as_uint8 packed

Available conversions:

• Floating point: as_float16, as_float32, as_float64
• Signed integers: as_int8, as_int16, as_int32, as_int64
• Unsigned integers: as_uint8, as_uint16
• Complex: as_complex32, as_complex64

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: reduce brightness and add bias *)
  let processed = Nx.add_s (Nx.mul_s normalized 0.8) 0.1 in

  (* Convert back to uint8 *)
  let result = Nx.astype uint8 (Nx.clip ~min:0.0 ~max:255.0 (Nx.mul_s processed 255.0)) in

  (* Save result *)
  save_image output_path result

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 path [
    ("epoch", P (Nx.scalar int32 epoch));
    ("weights", P weights);
    ("optimizer_state", P optimizer_state);
  ]

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

Error Handling

All I/O operations may raise Failure exceptions:

• File not found or inaccessible
• Unsupported file format
• Invalid data format
• Incompatible array shapes (for save_image)
• Missing archive members (for load_npz_member)

Performance Considerations

• Image loading/saving uses the stb_image libraries (header-only C libraries)
• NumPy format I/O is implemented in pure OCaml
• Large .npz archives are loaded entirely into memory
• For very large datasets, consider loading arrays individually with load_npz_member