package nx

  1. Overview
  2. Docs
package nx
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Module Nx_ioSource

Nx_io: Nx input/output for common file formats.

Provides functions to load and save Nx.t arrays in image formats, NumPy formats (.npy, .npz), HDF5 archives, and SafeTensors. Emphasizes safety with result types for error handling and labeled arguments for clarity.

Types

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

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

Wraps arrays whose element type and layout are determined at runtime, such as those loaded from files.

Sourcetype archive = (string, packed_nx) Hashtbl.t

Generic archive type for mappings from names/paths to packed arrays, used across NPZ, HDF5, and SafeTensors.

Conversions from Packed Arrays

Sourceval as_float16 : packed_nx -> Nx.float16_t

as_float16 packed converts a packed Nx to a Nx.float16_t, or Error if dtype mismatch.

Sourceval as_float32 : packed_nx -> Nx.float32_t

as_float32 packed converts a packed Nx to a Nx.float32_t, or Error if dtype mismatch.

Sourceval as_float64 : packed_nx -> Nx.float64_t

as_float64 packed converts a packed Nx to a Nx.float64_t, or Error if dtype mismatch.

Sourceval as_int8 : packed_nx -> Nx.int8_t

as_int8 packed converts a packed Nx to a Nx.int8_t, or Error if dtype mismatch.

Sourceval as_int16 : packed_nx -> Nx.int16_t

as_int16 packed converts a packed Nx to a Nx.int16_t, or Error if dtype mismatch.

Sourceval as_int32 : packed_nx -> Nx.int32_t

as_int32 packed converts a packed Nx to a Nx.int32_t, or Error if dtype mismatch.

Sourceval as_int64 : packed_nx -> Nx.int64_t

as_int64 packed converts a packed Nx to a Nx.int64_t, or Error if dtype mismatch.

Sourceval as_uint8 : packed_nx -> Nx.uint8_t

as_uint8 packed converts a packed Nx to a Nx.uint8_t, or Error if dtype mismatch.

Sourceval as_uint16 : packed_nx -> Nx.uint16_t

as_uint16 packed converts a packed Nx to a Nx.uint16_t, or Error if dtype mismatch.

Sourceval as_complex32 : packed_nx -> Nx.complex32_t

as_complex32 packed converts a packed Nx to a Nx.complex32_t, or Error if dtype mismatch.

Sourceval as_complex64 : packed_nx -> Nx.complex64_t

as_complex64 packed converts a packed Nx to a Nx.complex64_t, or Error if dtype mismatch.

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

  • Ok with uint8 nx of shape |height; width| if grayscale, or |height; width; 3| for RGB color images.
  • Error on failure (e.g., I/O or unsupported format).

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 *)
  match load_image "photo.png" with
  | Ok img -> (* img : (int, Nx.uint8_elt) Nx.t with shape [|H; W; 3|] *)
  | Error e -> failwith (string_of_error e)

  (* Load as grayscale *)
  let gray = load_image ~grayscale:true "photo.png"
Sourceval save_image : ?overwrite:bool -> string -> (int, Nx.uint8_elt) Nx.t -> unit

save_image ?overwrite path img

Save a uint8 nx as an image file.

Parameters

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

Returns

  • Ok () on success, Error on failure (e.g., unsupported shape or I/O).

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
  ignore (save_image ~img:gray "output.png")

  (* Save an RGB image without overwriting *)
  let rgb = Nx.create Nx.uint8 [|100; 200; 3|] data in
  save_image ~img:rgb "output.jpg" ~overwrite:false

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

  • Ok with packed_nx wrapping the loaded array.
  • Error on failure.

Examples

  (* Load and convert to specific type *)
  match load_npy "data.npy" with
  | Ok packed -> let arr = as_float32 packed in ...
  | Error _ -> ...
Sourceval save_npy : ?overwrite:bool -> string -> ('a, 'b) Nx.t -> unit

save_npy ~arr path ?overwrite

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

Parameters

  • ~arr: nx to save (any supported dtype).
  • path: destination path for the `.npy` file.
  • ?overwrite: if false, fail if file exists; default true.

Returns

  • Ok () on success, Error on failure.

Notes

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

NumPy Archive Format (.npz)

Sourceval load_npz : string -> archive

load_npz path

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

Returns

  • Ok with archive mapping each array name to its packed_nx.
Sourceval load_npz_member : name:string -> string -> packed_nx

load_npz_member ~name path

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

Returns

  • Ok with packed_nx or Error (e.g., Missing_entry).
Sourceval save_npz : ?overwrite:bool -> string -> (string * packed_nx) list -> unit

save_npz ?overwrite path items

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

SafeTensors Format

Sourceval load_safetensor : string -> archive

load_safetensor path

Load all tensors from a SafeTensors file into a hash table..

Sourceval save_safetensor : ?overwrite:bool -> string -> (string * packed_nx) list -> unit

save_safetensor ?overwrite path items

Save multiple named nxs to a SafeTensors file.

Sourcemodule Safe : sig ... end