Module Nx_ioSource

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.

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|] *)

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"

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

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"

Map from array names to packed_nx values for a NumPy `.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"

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"

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"

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

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

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

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

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

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

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

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

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

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

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