Module Nx_ioSource

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.

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

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

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

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

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

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

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

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

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

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

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

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

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"

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

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 _ -> ...

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.

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.

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).

save_npz ?overwrite path items

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

load_safetensor path

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

save_safetensor ?overwrite path items

Save multiple named nxs to a SafeTensors file.