package nx

  1. Overview
  2. Docs

Module Nx_ioSource

Nx_io: Nx input/output for common file formats.

Provide functions to load and save Nx.t arrays in image formats, NumPy formats (.npy, .npz), and HDF5 archives.

Sourcetype packed_nx =
  1. | P : ('a, 'b) Nx.t -> packed_nx

Existential container for an Nx.t of any dtype and dimensionality.

Wrap arrays whose element type and layout are determined at runtime, such as those loaded from .npy or .npz files.

Image Loading and Saving

Sourceval load_image : ?grayscale:bool -> string -> (int, Nx.uint8_elt) Nx.t

load_image ?grayscale path

Load an image file into a uint8 nx.

Parameters

  • ?grayscale: if true, load as grayscale (2D) else color (3D RGB); default false.
  • path: path to the image file.

Returns

  • uint8 nx of shape |height; width| if grayscale, or |height; width; 3| for RGB color images.

Raises

  • Failure if file I/O fails or format is unsupported.

Notes

  • Supported formats depend on stb_image (PNG, JPEG, BMP, TGA, GIF).
  • Alpha channels are discarded.
  • Pixel values range 0, 255.

Examples

  (* Load color image *)
  let img = load_image "photo.png" in
  (* img : (int, Nx.uint8_elt) Nx.t with shape [|H; W; 3|] *)

  (* Load as grayscale *)
  let gray = load_image ~grayscale:true "photo.png" in
  (* gray : (int, Nx.uint8_elt) Nx.t with shape [|H; W|] *)
Sourceval save_image : (int, Nx.uint8_elt) Nx.t -> string -> unit

save_image img path

Save a uint8 nx as an image file.

Parameters

  • img: uint8 nx of shape |height; width| (grayscale), |height; width; 1|, |height; width; 3| (RGB), or |height; width; 4| (RGBA).
  • path: destination file path; extension determines format (e.g. .png, .jpg).

Raises

  • Failure if nx is not uint8 kind or has unsupported dimensions.
  • Failure on I/O error.

Notes

  • Supported formats depend on stb_image_write library.
  • Pixel values are written in 0, 255 range.

Examples

  (* Save a grayscale image *)
  let gray = Nx.create Nx.uint8 [|100; 200|] data in
  save_image gray "output.png"

  (* Save an RGB image *)
  let rgb = Nx.create Nx.uint8 [|100; 200; 3|] data in
  save_image rgb "output.jpg"

NumPy Format (.npy)

Sourceval load_npy : string -> packed_nx

load_npy path

Load a single nx from a NumPy `.npy` file.

Parameters

  • path: path to the `.npy` file.

Returns

  • packed_nx wrapping the loaded array with its runtime-detected dtype.

Raises

  • Failure if file I/O fails or format is invalid.

Examples

  (* Load and convert to specific type *)
  let arr = load_npy "data.npy" |> to_float32 in

  (* Pattern match to unpack *)
  let P arr = load_npy "data.npy" in
  (* arr : ('a, 'b) Nx.t *)
Sourceval save_npy : ('a, 'b) Nx.t -> string -> unit

save_npy arr path

Save a single nx to a NumPy `.npy` file.

Parameters

  • arr: nx to save (any supported dtype).
  • path: destination path for the `.npy` file.

Raises

  • Failure on I/O error or unsupported dtype.

Notes

  • The file encodes dtype, shape, and raw data in little-endian order.

Examples

  let arr = Nx.arange Nx.float32 0.0 10.0 1.0 in
  save_npy arr "array.npy"

NumPy Archive Format (.npz)

Sourcetype npz_archive = (string, packed_nx) Hashtbl.t

Map from array names to packed_nx values for a NumPy `.npz` archive.

Sourceval load_npz : string -> npz_archive

load_npz path

Load all arrays from a NumPy `.npz` archive into a hash table.

Parameters

  • path: path to the `.npz` file.

Returns

  • npz_archive mapping each array name to its packed_nx.

Raises

  • Failure on I/O error or invalid archive format.

Examples

  let archive = load_npz "bundle.npz" in
  match Hashtbl.find_opt archive "weights" with
  | Some (P weights) ->
      let w = to_float32 (P weights) in
      (* use weights *)
  | None -> failwith "weights not found"
Sourceval load_npz_member : path:string -> name:string -> packed_nx

load_npz_member ~path ~name

Load a single named array from a NumPy `.npz` archive.

Parameters

  • path: path to the `.npz` file.
  • name: name of the array entry in the archive.

Returns

  • packed_nx containing the loaded array.

Raises

  • Failure if entry name is not found or on I/O error.

Examples

  (* Load specific array without loading entire archive *)
  let (P data) = load_npz_member ~path:"model.npz" ~name:"embeddings"
Sourceval save_npz : (string * packed_nx) list -> string -> unit

save_npz items path

Save multiple named nxs to a NumPy `.npz` archive.

Parameters

  • items: list of (name, array) pairs to include in the archive.
  • path: destination path for the `.npz` file.

Raises

  • Failure on I/O error or archive creation failure.

Examples

  (* Save multiple arrays *)
  let weights = Nx.randn Nx.float32 [| 784; 128 |] in
  let bias = Nx.zeros Nx.float32 [| 128 |] in
  save_npz [ ("weights", P weights); ("bias", P bias) ] "model.npz"

Conversions from Packed Arrays

Sourceval to_float16 : packed_nx -> Nx.float16_t

to_float16 packed_nx converts a packed Nx to a Nx.float16_t.

Sourceval to_float32 : packed_nx -> Nx.float32_t

to_float32 packed_nx converts a packed Nx to a Nx.float32_t.

Sourceval to_float64 : packed_nx -> Nx.float64_t

to_float64 packed_nx converts a packed Nx to a Nx.float64_t.

Sourceval to_int8 : packed_nx -> Nx.int8_t

to_int8 packed_nx converts a packed Nx to a Nx.int8_t.

Sourceval to_int16 : packed_nx -> Nx.int16_t

to_int16 packed_nx converts a packed Nx to a Nx.int16_t.

Sourceval to_int32 : packed_nx -> Nx.int32_t

to_int32 packed_nx converts a packed Nx to a Nx.int32_t.

Sourceval to_int64 : packed_nx -> Nx.int64_t

to_int64 packed_nx converts a packed Nx to a Nx.int64_t.

Sourceval to_uint8 : packed_nx -> Nx.uint8_t

to_uint8 packed_nx converts a packed Nx to a Nx.uint8_t.

Sourceval to_uint16 : packed_nx -> Nx.uint16_t

to_uint16 packed_nx converts a packed Nx to a Nx.uint16_t.

Sourceval to_complex32 : packed_nx -> Nx.complex32_t

to_complex32 packed_nx converts a packed Nx to a Nx.complex32_t.

Sourceval to_complex64 : packed_nx -> Nx.complex64_t

to_complex64 packed_nx converts a packed Nx to a Nx.complex64_t.

OCaml

Innovation. Community. Security.