01-creating-arrays

Build arrays from scratch — constants, ranges, grids, and custom data. This example walks through the most common ways to create arrays in Nx.

dune exec nx/examples/01-creating-arrays/main.exe

What You'll Learn

  • Choosing a dtype (float32, float64, int32)
  • Filling arrays with constants: zeros, ones, full
  • Generating ranges: arange, linspace, logspace
  • Building arrays from OCaml data: create, init
  • Diagonal and special matrices: identity, eye, tril, triu
  • Coordinate grids with meshgrid

Key Functions

Function Purpose
zeros dtype shape Array of all zeros
ones dtype shape Array of all ones
full dtype shape value Array filled with a value
arange dtype start stop step Integer-stepped range (exclusive stop)
linspace dtype start stop n Evenly spaced floats
logspace dtype start stop n Logarithmically spaced values
create dtype shape data Array from an OCaml array
init dtype shape f Array from a function of indices
identity dtype n n×n identity matrix
eye ?k dtype n Ones on the k-th diagonal
meshgrid x y Coordinate grids from 1D arrays
tril m / triu m Lower / upper triangular part

Output Walkthrough

When you run this example, you'll see arrays printed in a compact format:

zeros (2×3):
[[0, 0, 0],
 [0, 0, 0]]

arange 0..9:
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

5×5 multiplication table:
[[1, 2, 3, 4, 5],
 [2, 4, 6, 8, 10],
 [3, 6, 9, 12, 15],
 [4, 8, 12, 16, 20],
 [5, 10, 15, 20, 25]]

The multiplication table is built with init, which calls a function with the index array for each element:

init int32 [| 5; 5 |] (fun idx ->
    Int32.of_int ((idx.(0) + 1) * (idx.(1) + 1)))

meshgrid builds a pair of 2D coordinate grids from two 1D arrays — useful for evaluating functions over a grid:

meshgrid X:
[[0, 1, 2],
 [0, 1, 2]]
meshgrid Y:
[[0, 0, 0],
 [1, 1, 1]]

Dtypes

Every array has a dtype that determines the element type and precision. The first argument to most creation functions is the dtype:

zeros float32 [| 2; 3 |]   (* 32-bit floats *)
ones float64 [| 3 |]       (* 64-bit floats *)
arange int32 0 10 1        (* 32-bit integers *)

Nx supports 18 dtypes including Float16, BFloat16, Complex128, Bool, and various integer widths.

Try It

  1. Create a 10-element linspace from -1.0 to 1.0 and print it.
  2. Use init to build a 4×4 matrix where each element is the sum of its row and column index.
  3. Try eye ~k:(-1) float64 4 to see the subdiagonal.

Next Steps

Continue to 02-infix-and-arithmetic to learn how the Infix module makes array math read like algebra.

(** Build arrays from scratch — constants, ranges, grids, and custom data.

    This example walks through the most common ways to create arrays. By the end
    you'll know how to pick a dtype, fill arrays with constants, generate
    ranges, and build grids and triangular matrices. *)

open Nx

let () =
  (* Constant-filled arrays: zeros, ones, and an arbitrary fill value. *)
  let z = zeros float32 [| 2; 3 |] in
  Printf.printf "zeros (2×3):\n%s\n\n" (data_to_string z);

  let o = ones float64 [| 3 |] in
  Printf.printf "ones (3):\n%s\n\n" (data_to_string o);

  let pi = full float64 [| 2; 2 |] Float.pi in
  Printf.printf "full π (2×2):\n%s\n\n" (data_to_string pi);

  (* Ranges: integer steps and evenly-spaced floats. *)
  let ints = arange int32 0 10 1 in
  Printf.printf "arange 0..9:\n%s\n\n" (data_to_string ints);

  let spaced = linspace float64 0.0 1.0 5 in
  Printf.printf "linspace 0..1 (5 points):\n%s\n\n" (data_to_string spaced);

  let decades = logspace float64 1.0 4.0 4 in
  Printf.printf "logspace 10¹..10⁴:\n%s\n\n" (data_to_string decades);

  (* From raw data: pack an OCaml array into a 2×3 matrix. *)
  let data = create float64 [| 2; 3 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in
  Printf.printf "create from data (2×3):\n%s\n\n" (data_to_string data);

  (* Build a multiplication table with [init]. *)
  let mul_table =
    init int32 [| 5; 5 |] (fun idx ->
        Int32.of_int ((idx.(0) + 1) * (idx.(1) + 1)))
  in
  Printf.printf "5×5 multiplication table:\n%s\n\n" (data_to_string mul_table);

  (* Identity and eye: diagonal matrices. *)
  let id = identity float64 3 in
  Printf.printf "identity 3×3:\n%s\n\n" (data_to_string id);

  let e = eye ~k:1 float64 3 in
  Printf.printf "eye (k=1, superdiagonal):\n%s\n\n" (data_to_string e);

  (* Coordinate grids with meshgrid. *)
  let xs = arange_f float64 0.0 3.0 1.0 in
  let ys = arange_f float64 0.0 2.0 1.0 in
  let grid_x, grid_y = meshgrid xs ys in
  Printf.printf "meshgrid X:\n%s\n" (data_to_string grid_x);
  Printf.printf "meshgrid Y:\n%s\n\n" (data_to_string grid_y);

  (* Triangular matrices: tril and triu. *)
  let m = ones float64 [| 4; 4 |] in
  Printf.printf "tril (lower triangle):\n%s\n" (data_to_string (tril m));
  Printf.printf "triu (upper triangle):\n%s\n" (data_to_string (triu m))