Module Nx

The type for tensors with OCaml element type 'a and buffer element kind 'b.

The type for index specifications used by slice and set_slice.

data t is the underlying flat buffer of t.

The buffer is shared: mutations through the buffer are visible through t and vice-versa. The buffer may be larger than the tensor's logical extent when t is a strided view.

shape t is the dimensions of t. A scalar tensor has shape |\||.

dtype t is the data type of t.

strides t is the byte stride for each dimension of t.

Raises Invalid_argument if t does not have computable strides (e.g. after certain non-contiguous view operations). Use is_c_contiguous or call contiguous first.

See also stride.

stride i t is the byte stride of dimension i.

Raises Invalid_argument if i is out of bounds or t does not have computable strides.

See also strides.

dims t is Shape manipulation.

dim i t is the size of dimension i.

Raises Invalid_argument if i is out of bounds.

ndim t is the number of dimensions of t.

itemsize t is the number of bytes per element.

size t is the total number of elements.

numel t is size.

nbytes t is size t * itemsize t.

offset t is the element offset of t in its underlying buffer.

is_c_contiguous t is true iff t's elements are laid out contiguously in row-major (C) order.

See also contiguous.

to_bigarray t is a contiguous bigarray with the same shape and data as t. Always copies.

Raises Invalid_argument if t's dtype is an extended type not supported by Bigarray.

See also of_bigarray.

to_buffer t is a flat, contiguous buffer of t's data.

Returns the underlying buffer directly when t is already contiguous with zero offset and matching size; copies otherwise.

to_array t is a fresh OCaml array containing the elements of t in row-major order. Always copies.

  # let t =
      create int32 [| 2; 2 |] [| 1l; 2l; 3l; 4l |]
    in
    to_array t
  - : int32 array = [|1l; 2l; 3l; 4l|]

create dtype shape data is a tensor of the given dtype and shape initialised from data in row-major order.

Raises Invalid_argument if Array.length data does not equal the product of shape.

  # create float32 [| 2; 3 |]
      [| 1.; 2.; 3.; 4.; 5.; 6. |]
  - : (float, float32_elt) t = float32 [2; 3] [[1, 2, 3],
                                               [4, 5, 6]]

init dtype shape f is a tensor where the element at multi-index i is f i.

  # init int32 [| 2; 3 |]
      (fun i -> Int32.of_int (i.(0) + i.(1)))
  - : (int32, int32_elt) t = int32 [2; 3] [[0, 1, 2],
                                           [1, 2, 3]]

empty dtype shape is an uninitialized tensor.

Warning. Elements contain arbitrary values until written.

full dtype shape v is a tensor filled with v.

  # full float32 [| 2; 3 |] 3.14
  - : (float, float32_elt) t = float32 [2; 3]
  [[3.14, 3.14, 3.14],
   [3.14, 3.14, 3.14]]

ones dtype shape is a tensor filled with ones.

zeros dtype shape is a tensor filled with zeros.

scalar dtype v is a 0-dimensional tensor containing v. The result has shape |\||.

empty_like t is empty with the same dtype and shape as t.

full_like t v is full with the same dtype and shape as t.

ones_like t is ones with the same dtype and shape as t.

zeros_like t is zeros with the same dtype and shape as t.

scalar_like t v is scalar with the same dtype as t.

eye ?m ?k dtype n is an n × m matrix with ones on the k-th diagonal and zeros elsewhere. m defaults to n. k defaults to 0 (main diagonal); positive k selects an upper diagonal, negative k a lower one.

  # eye int32 3
  - : (int32, int32_elt) t = int32 [3; 3] [[1, 0, 0],
                                           [0, 1, 0],
                                           [0, 0, 1]]
  # eye ~k:1 int32 3
  - : (int32, int32_elt) t = int32 [3; 3] [[0, 1, 0],
                                           [0, 0, 1],
                                           [0, 0, 0]]

See also identity, diag.

identity dtype n is eye dtype n.

diag ?k v extracts or constructs a diagonal.

When v is 1-D, returns a 2-D tensor with v on the k-th diagonal. When v is 2-D, returns the k-th diagonal as a 1-D tensor. k defaults to 0.

Raises Invalid_argument if v is not 1-D or 2-D.

  # let v = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    diag v
  - : (int32, int32_elt) t = int32 [3; 3] [[1, 0, 0],
                                           [0, 2, 0],
                                           [0, 0, 3]]
  # let x =
      arange int32 0 9 1 |> reshape [| 3; 3 |]
    in
    diag x
  - : (int32, int32_elt) t = [0, 4, 8]

See also eye, diagonal.

arange dtype start stop step is a 1-D tensor of values from start (inclusive) to stop (exclusive) with stride step.

Raises Invalid_argument if step = 0.

  # arange int32 0 10 2
  - : (int32, int32_elt) t = int32 [5] [0, 2, ..., 6, 8]
  # arange int32 5 0 (-1)
  - : (int32, int32_elt) t = int32 [5] [5, 4, ..., 2, 1]

See also arange_f, linspace.

arange_f dtype start stop step is like arange for floating-point ranges.

Raises Invalid_argument if step = 0.0.

  # arange_f float32 0. 1. 0.2
  - : (float, float32_elt) t = float32 [5] [0, 0.2, ..., 0.6, 0.8]

See also arange, linspace.

linspace dtype ?endpoint start stop n is n values evenly spaced from start to stop. endpoint defaults to true (include stop).

Raises Invalid_argument if n is negative.

  # linspace float32 0. 10. 5
  - : (float, float32_elt) t = float32 [5] [0, 2.5, ..., 7.5, 10]
  # linspace float32 ~endpoint:false 0. 10. 5
  - : (float, float32_elt) t = float32 [5] [0, 2, ..., 6, 8]

See also logspace, geomspace.

logspace dtype ?endpoint ?base start stop n is n values evenly spaced on a logarithmic scale: base{^x} where x ranges from start to stop. endpoint defaults to true. base defaults to 10.0.

Raises Invalid_argument if n is negative.

  # logspace float32 0. 2. 3
  - : (float, float32_elt) t = [1, 10, 100]
  # logspace float32 ~base:2.0 0. 3. 4
  - : (float, float32_elt) t = [1, 2, 4, 8]

See also linspace, geomspace.

geomspace dtype ?endpoint start stop n is n values evenly spaced on a geometric (multiplicative) scale. endpoint defaults to true.

Raises Invalid_argument if start or stop is not positive.

  # geomspace float32 1. 1000. 4
  - : (float, float32_elt) t = [1, 10, 100, 1000]

See also linspace, logspace.

meshgrid ?indexing x y is a pair of 2-D coordinate grids built from 1-D arrays x and y. indexing defaults to `xy (Cartesian: X varies along columns, Y along rows). With `ij (matrix), X varies along rows, Y along columns.

Raises Invalid_argument if x or y is not 1-D.

  # let x = linspace float32 0. 2. 3 in
    let y = linspace float32 0. 1. 2 in
    meshgrid x y
  - : (float, float32_elt) t * (float, float32_elt) t =
  (float32 [2; 3] [[0, 1, 2],
                   [0, 1, 2]], float32 [2; 3] [[0, 0, 0],
                                               [1, 1, 1]])

tril ?k x is the lower-triangular part of x with elements above the k-th diagonal set to zero. k defaults to 0.

Raises Invalid_argument if x has fewer than 2 dimensions.

See also triu.

triu ?k x is the upper-triangular part of x with elements below the k-th diagonal set to zero. k defaults to 0.

Raises Invalid_argument if x has fewer than 2 dimensions.

See also tril.

of_bigarray ba is a tensor sharing memory with ba.

Zero-copy: mutations through either are visible to both.

See also to_bigarray.

of_buffer buf ~shape is a tensor viewing buf with the given shape. The product of shape must equal the buffer length.

one_hot ~num_classes indices is a one-hot encoded tensor.

Appends a new trailing dimension of size num_classes. Values in indices must lie in [0, num_classes). Out-of-range indices produce all-zero rows.

Raises Invalid_argument if indices is not an integer dtype or num_classes <= 0.

  # let idx =
      create int32 [| 3 |] [| 0l; 1l; 3l |]
    in
    one_hot ~num_classes:4 idx
  - : (int, uint8_elt) t = uint8 [3; 4]
  [[1, 0, 0, 0],
   [0, 1, 0, 0],
   [0, 0, 0, 1]]

Parameter trees: structures with tensor leaves. Ptree.S is the traversal interface shared across the Raven ecosystem; Ptree.t is the stock dynamic instance.

Splittable RNG keys and implicit key management.

rand dtype shape samples uniformly from [0, 1).

Raises Invalid_argument if dtype is not a float type.

randn dtype shape samples from the standard normal distribution (mean 0, variance 1) via the Box–Muller transform.

Raises Invalid_argument if dtype is not a float type.

randint dtype ?high shape low samples integers uniformly from [low, high). high defaults to 10.

Raises Invalid_argument if dtype is not an integer type or low >= high.

bernoulli ~p shape samples booleans that are true with probability p.

Raises Invalid_argument if p is not in [0, 1].

permutation n is a random permutation of [0, n-1].

Raises Invalid_argument if n <= 0.

shuffle t is a copy of t with the first axis randomly permuted. No-op on scalars.

categorical ?axis ?shape logits samples category indices from unnormalised log-probabilities using the Gumbel-max trick. axis defaults to -1 (last axis). shape prepends extra batch dimensions.

Raises Invalid_argument if logits is not a float type or axis is out of bounds.

truncated_normal dtype ~lower ~upper shape samples from a standard normal distribution truncated to [lower, upper].

Raises Invalid_argument if dtype is not a float type or lower >= upper.

reshape shape t is a view of t with the given shape.

At most one dimension may be -1; it is inferred from the total number of elements. The product of shape must equal size t.

Raises Invalid_argument if shape is incompatible or contains more than one -1.

  # create int32 [| 6 |] [| 1l; 2l; 3l; 4l; 5l; 6l |]
    |> reshape [| 2; 3 |]
  - : (int32, int32_elt) t = int32 [2; 3] [[1, 2, 3],
                                           [4, 5, 6]]
  # create int32 [| 6 |] [| 1l; 2l; 3l; 4l; 5l; 6l |]
    |> reshape [| 3; -1 |]
  - : (int32, int32_elt) t = int32 [3; 2] [[1, 2],
                                           [3, 4],
                                           [5, 6]]

See also flatten, unflatten, ravel.

broadcast_to shape t is a view of t broadcast to shape.

Dimensions are aligned from the right; each dimension of t must be 1 or equal to the corresponding target dimension. Broadcast dimensions have zero byte-stride (no copy).

Raises Invalid_argument if the shapes are incompatible.

  # create int32 [| 1; 3 |] [| 1l; 2l; 3l |]
    |> broadcast_to [| 3; 3 |]
  - : (int32, int32_elt) t = int32 [3; 3] [[1, 2, 3],
                                           [1, 2, 3],
                                           [1, 2, 3]]

See also broadcasted, expand.

broadcasted ?reverse t1 t2 is (t1', t2') where both are broadcast to their common shape. When reverse is true (default false), returns (t2', t1').

Raises Invalid_argument if the shapes are incompatible.

See also broadcast_to, broadcast_arrays.

expand shape t is like broadcast_to but -1 in shape preserves the corresponding dimension of t.

Raises Invalid_argument if any dimension in shape is negative (other than -1).

  # ones float32 [| 1; 4; 1 |]
    |> expand [| 3; -1; 5 |] |> shape
  - : int array = [|3; 4; 5|]

See also broadcast_to.

flatten ?start_dim ?end_dim t collapses dimensions start_dim through end_dim (inclusive) into a single dimension. start_dim defaults to 0. end_dim defaults to -1 (last). Negative indices count from the end.

Raises Invalid_argument if indices are out of bounds.

  # zeros float32 [| 2; 3; 4 |] |> flatten |> shape
  - : int array = [|24|]
  # zeros float32 [| 2; 3; 4; 5 |]
    |> flatten ~start_dim:1 ~end_dim:2 |> shape
  - : int array = [|2; 12; 5|]

See also unflatten, ravel.

unflatten dim sizes t expands dimension dim into multiple dimensions given by sizes. At most one element of sizes may be -1 (inferred). The product of sizes must equal the size of dimension dim.

Raises Invalid_argument if the product mismatches or dim is out of bounds.

  # zeros float32 [| 2; 12; 5 |]
    |> unflatten 1 [| 3; 4 |] |> shape
  - : int array = [|2; 3; 4; 5|]

See also flatten.

ravel t is t reshaped to 1-D. Returns a view when possible.

Raises Invalid_argument if t cannot be flattened without copying; call contiguous first.

See also flatten, contiguous.

squeeze ?axes t removes dimensions of size 1. When axes is given, only those axes are removed. Negative indices count from the end.

Raises Invalid_argument if a specified axis does not have size 1.

  # ones float32 [| 1; 3; 1; 4 |]
    |> squeeze |> shape
  - : int array = [|3; 4|]
  # ones float32 [| 1; 3; 1; 4 |]
    |> squeeze ~axes:[ 0 ] |> shape
  - : int array = [|3; 1; 4|]

See also unsqueeze.

unsqueeze ?axes t inserts dimensions of size 1 at the positions listed in axes. Positions refer to the result tensor.

Raises Invalid_argument if axes is not specified, contains duplicates, or values are out of bounds.

  # create float32 [| 3 |] [| 1.; 2.; 3. |]
    |> unsqueeze ~axes:[ 0; 2 ] |> shape
  - : int array = [|1; 3; 1|]

See also squeeze, expand_dims.

squeeze_axis i t removes dimension i if its size is 1.

Raises Invalid_argument if dimension i is not 1.

See also squeeze.

unsqueeze_axis i t inserts a dimension of size 1 at position i.

See also unsqueeze.

expand_dims axes t is unsqueeze ~axes t.

transpose ?axes t permutes the dimensions of t.

axes must be a permutation of [0; …; ndim t - 1]. When omitted, reverses all dimensions. Returns a view (no copy).

Raises Invalid_argument if axes is not a valid permutation.

  # create int32 [| 2; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l |]
    |> transpose
  - : (int32, int32_elt) t = int32 [3; 2] [[1, 4],
                                           [2, 5],
                                           [3, 6]]

See also matrix_transpose, moveaxis, swapaxes.

flip ?axes t reverses elements along the given axes. When omitted, flips all dimensions.

Raises Invalid_argument if any axis is out of bounds.

  # create int32 [| 2; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l |]
    |> flip ~axes:[ 1 ]
  - : (int32, int32_elt) t = int32 [2; 3] [[3, 2, 1],
                                           [6, 5, 4]]

moveaxis src dst t moves dimension src to position dst.

Raises Invalid_argument if either index is out of bounds.

See also transpose, swapaxes.

swapaxes a1 a2 t exchanges dimensions a1 and a2.

Raises Invalid_argument if either index is out of bounds.

See also transpose, moveaxis.

roll ?axis shift t shifts elements along axis by shift positions, wrapping around. When axis is omitted, operates on the flattened tensor. Negative shift rolls backward.

Raises Invalid_argument if axis is out of bounds.

  # create int32 [| 5 |] [| 1l; 2l; 3l; 4l; 5l |]
    |> roll 2
  - : (int32, int32_elt) t = int32 [5] [4, 5, ..., 2, 3]

pad widths value t pads t with value. widths.(i) is (before, after) for dimension i.

Raises Invalid_argument if Array.length widths does not match ndim t or any width is negative.

  # create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |]
    |> pad [| (1, 1); (1, 1) |] 0. |> shape
  - : int array = [|4; 4|]

See also shrink.

shrink ranges t extracts a slice where ranges.(i) is (start, stop) (exclusive) for dimension i. Returns a view.

  # create int32 [| 3; 3 |]
      [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |]
    |> shrink [| (1, 3); (0, 2) |]
  - : (int32, int32_elt) t = int32 [2; 2] [[4, 5],
                                           [7, 8]]

See also pad.

tile reps t is t repeated according to reps. reps.(i) gives the repetition count along dimension i. If reps is longer than ndim t, dimensions are prepended.

Raises Invalid_argument if any repetition count is negative.

  # create int32 [| 1; 2 |] [| 1l; 2l |]
    |> tile [| 2; 3 |]
  - : (int32, int32_elt) t = int32 [2; 6] [[1, 2, ..., 1, 2],
                                           [1, 2, ..., 1, 2]]

See also repeat.

repeat ?axis n t repeats each element n times along axis. When axis is omitted, operates on the flattened tensor.

Raises Invalid_argument if n is negative or axis is out of bounds.

  # create int32 [| 3 |] [| 1l; 2l; 3l |]
    |> repeat 2
  - : (int32, int32_elt) t = int32 [6] [1, 1, ..., 3, 3]

See also tile.

concatenate ?axis ts joins tensors along an existing axis. All tensors must have the same shape except on the concatenation axis. When axis is omitted, every tensor is flattened first. Always copies.

Raises Invalid_argument if the list is empty or shapes are incompatible.

  # let a = create int32 [| 2; 2 |] [| 1l; 2l; 3l; 4l |] in
    let b = create int32 [| 1; 2 |] [| 5l; 6l |] in
    concatenate ~axis:0 [ a; b ]
  - : (int32, int32_elt) t = int32 [3; 2] [[1, 2],
                                           [3, 4],
                                           [5, 6]]

See also stack, vstack, hstack.

stack ?axis ts joins tensors along a new axis. All tensors must have identical shape. axis defaults to 0. Negative values count from the end of the result shape.

Raises Invalid_argument if the list is empty, shapes differ, or axis is out of bounds.

  # let a = create int32 [| 2 |] [| 1l; 2l |] in
    let b = create int32 [| 2 |] [| 3l; 4l |] in
    stack [ a; b ]
  - : (int32, int32_elt) t = int32 [2; 2] [[1, 2],
                                           [3, 4]]
  # let a = create int32 [| 2 |] [| 1l; 2l |] in
    let b = create int32 [| 2 |] [| 3l; 4l |] in
    stack ~axis:1 [ a; b ]
  - : (int32, int32_elt) t = int32 [2; 2] [[1, 3],
                                           [2, 4]]

See also concatenate.

vstack ts stacks vertically (along axis 0). 1-D tensors are treated as row vectors (shape [1; n]).

Raises Invalid_argument if shapes are incompatible.

  # let a = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    let b = create int32 [| 3 |] [| 4l; 5l; 6l |] in
    vstack [ a; b ]
  - : (int32, int32_elt) t = int32 [2; 3] [[1, 2, 3],
                                           [4, 5, 6]]

See also hstack, dstack, concatenate.

hstack ts stacks horizontally. 1-D tensors are concatenated directly; higher-D tensors concatenate along axis 1.

Raises Invalid_argument if shapes are incompatible.

  # let a = create int32 [| 2; 1 |] [| 1l; 2l |] in
    let b = create int32 [| 2; 1 |] [| 3l; 4l |] in
    hstack [ a; b ]
  - : (int32, int32_elt) t = int32 [2; 2] [[1, 3],
                                           [2, 4]]

See also vstack, dstack, concatenate.

dstack ts stacks depth-wise (along axis 2). Tensors are reshaped to at least 3-D before concatenation: 1-D [n] → [1; n; 1], 2-D [m; n] → [m; n; 1].

Raises Invalid_argument if the resulting shapes are incompatible.

See also vstack, hstack, concatenate.

broadcast_arrays ts broadcasts every tensor to their common shape. Returns views (no copies).

Raises Invalid_argument if shapes are incompatible.

See also broadcast_to, broadcasted.

array_split ~axis spec t splits t into sub-tensors.

With `Count n, divides as evenly as possible (first sections absorb extra elements). With `Indices [i0; i1; …], splits at the given indices producing [0, i0), [i0, i1), …, [ik, end).

Raises Invalid_argument if axis is out of bounds or spec is invalid.

  # create int32 [| 5 |] [| 1l; 2l; 3l; 4l; 5l |]
    |> array_split ~axis:0 (`Count 3)
  - : (int32, int32_elt) t list = [[1, 2]; [3, 4]; [5]]

See also split.

split ~axis n t splits t into n equal parts along axis.

Raises Invalid_argument if the axis size is not divisible by n.

See also array_split.

cast dtype t is a copy of t with elements converted to dtype.

  # create float32 [| 3 |] [| 1.5; 2.7; 3.1 |]
    |> cast int32
  - : (int32, int32_elt) t = [1, 2, 3]

See also contiguous, copy.

astype dtype t is cast.

contiguous t is t if it is already C-contiguous, or a fresh contiguous copy otherwise.

See also is_c_contiguous, copy.

copy t is a deep copy of t. Always allocates new memory; the result is contiguous.

  # let x = create float32 [| 3 |] [| 1.; 2.; 3. |] in
    let y = copy x in
    set_item [ 0 ] 999. y;
    x, y
  - : (float, float32_elt) t * (float, float32_elt) t =
  ([1, 2, 3], [999, 2, 3])

See also contiguous.

blit src dst copies the elements of src into dst in-place. Shapes must match exactly.

Raises Invalid_argument if shapes differ.

fill v t is a fresh copy of t with every element set to v. Does not mutate t.

get indices t is the sub-tensor at indices, indexing from the outermost dimension inward. Returns a scalar tensor when all dimensions are indexed; otherwise a view of the remaining dimensions. Negative indices count from the end.

Raises Invalid_argument if any index is out of bounds.

  # let x =
      create int32 [| 2; 3 |]
        [| 1l; 2l; 3l; 4l; 5l; 6l |]
    in
    get [ 1 ] x
  - : (int32, int32_elt) t = [4, 5, 6]

See also item, slice.

set indices t v writes v at the position given by indices.

Raises Invalid_argument if indices are out of bounds.

slice specs t extracts a sub-tensor using advanced indexing.

Each element of specs addresses one axis from left to right:

• I i — single index (reduces dimension; negative from end).
• L [i0; i1; …] — gather listed indices.
• R (start, stop) — half-open range [start, stop).
• Rs (start, stop, step) — strided range.
• A — full axis (default for trailing axes).
• M mask — boolean mask selecting positions where mask is true.
• N — insert a new axis of size 1.

Returns a view when possible.

Raises Invalid_argument if specs are out of bounds, if step is zero, or if a mask spec is used (not yet supported).

  # let x =
      create int32 [| 3; 3 |]
        [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |]
    in
    slice [ R (0, 2); L [ 0; 2 ] ] x
  - : (int32, int32_elt) t = int32 [2; 2] [[1, 3],
                                           [4, 6]]

See also get, set_slice.

set_slice specs t v writes v into the region of t selected by specs. v is broadcast if needed.

Raises Invalid_argument if N (new-axis) specs are used (not supported for writes).

See also slice.

item indices t is the scalar value at indices. Indices must cover all dimensions.

Raises Invalid_argument if the number of indices is wrong or any index is out of bounds.

See also get, set_item.

set_item indices v t sets the element at indices to v in-place. Indices must cover all dimensions.

Raises Invalid_argument if the number of indices is wrong or any index is out of bounds.

See also item.

take ?axis ?mode indices t gathers elements from t at indices along axis. When axis is omitted, t is flattened first. mode controls out-of-bounds indices: `raise (default) raises, `wrap uses modular indexing, `clip clamps to bounds.

Raises Invalid_argument if mode is `raise and any index is out of bounds.

  # let x =
      create int32 [| 5 |]
        [| 0l; 1l; 2l; 3l; 4l |]
    in
    take
      (create int32 [| 3 |] [| 1l; 3l; 0l |])
      x
  - : (int32, int32_elt) t = [1, 3, 0]

See also put, take_along_axis.

take_along_axis ~axis indices t gathers values from t along axis using indices. indices must match t's shape except along axis. Useful for gathering from argmax/argmin results.

Raises Invalid_argument if shapes are incompatible.

  # let x =
      create float32 [| 2; 3 |]
        [| 4.; 1.; 2.; 3.; 5.; 6. |]
    in
    let idx =
      create int32 [| 2; 1 |] [| 1l; 0l |]
    in
    take_along_axis ~axis:1 idx x
  - : (float, float32_elt) t = float32 [2; 1] [[1],
                                               [3]]

See also take, put_along_axis.

put ?axis ~indices ~values ?mode t writes values into t at positions given by indices. When axis is omitted, t is flattened first. mode defaults to `raise. Modifies t in-place.

Raises Invalid_argument if mode is `raise and any index is out of bounds.

See also take, put_along_axis, index_put.

index_put ~indices ~values ?mode t writes values into t at the coordinates given by indices.

indices contains one index tensor per axis of t; they are broadcast to a common shape that determines the number of updates. values is broadcast to the same shape. Duplicate coordinates overwrite. mode defaults to `raise.

Raises Invalid_argument if the number of index tensors does not match ndim t.

  # let t = zeros float32 [| 3; 3 |] in
    let rows =
      create int32 [| 3 |] [| 0l; 2l; 1l |]
    in
    let cols =
      create int32 [| 3 |] [| 1l; 0l; 2l |]
    in
    index_put ~indices:[| rows; cols |]
      ~values:(create float32 [| 3 |]
                 [| 10.; 20.; 30. |])
      t;
    t
  - : (float, float32_elt) t = float32 [3; 3]
  [[0, 10, 0],
   [0, 0, 30],
   [20, 0, 0]]

See also put.

put_along_axis ~axis ~indices ~values t writes values into t at positions selected by indices along axis. Modifies t in-place.

Raises Invalid_argument if shapes are incompatible.

See also take_along_axis, put, scatter.

scatter ?mode ?unique_indices ~axis ~indices ~values t is t with values placed at the positions selected by indices along axis: the pure counterpart of put_along_axis. indices must match t's shape except along axis, and values is broadcast to indices' shape. Index values must be in range for t's size along axis.

mode controls how updates combine with t: `Set (default) overwrites, the last update winning at duplicate positions; `Add accumulates every update into t's value. unique_indices = true promises that no position is selected twice, letting backends skip duplicate handling; the result is undefined if the promise is broken.

Unlike the in-place writes, scatter stays a pure tensor expression, so it traces under Rune.jit and differentiates with respect to both t and values.

  # let x = zeros float32 [| 2; 3 |] in
    let idx =
      create int32 [| 2; 1 |] [| 1l; 0l |]
    in
    scatter ~axis:1 ~indices:idx
      ~values:(create float32 [| 2; 1 |]
                 [| 10.; 20. |])
      x
  - : (float, float32_elt) t = float32 [2; 3] [[0, 10, 0],
                                               [20, 0, 0]]

Raises Invalid_argument if shapes are incompatible.

See also put_along_axis, take_along_axis.

compress ?axis ~condition t selects elements where condition is true along axis. condition must be 1-D. When axis is omitted, t is flattened first.

Raises Invalid_argument if the condition length is incompatible.

  # let x =
      create int32 [| 5 |]
        [| 1l; 2l; 3l; 4l; 5l |]
    in
    compress
      ~condition:(create bool [| 5 |]
        [| true; false; true; false; true |])
      x
  - : (int32, int32_elt) t = [1, 3, 5]

See also extract, nonzero.

extract ~condition t is the 1-D tensor of elements of t where condition is true. Both are flattened before comparison.

Raises Invalid_argument if sizes differ.

See also compress, nonzero.

nonzero t is an array of 1-D index tensors, one per dimension, giving the coordinates of non-zero elements.

  # let x =
      create int32 [| 3; 3 |]
        [| 0l; 1l; 0l;
           2l; 0l; 3l;
           0l; 0l; 4l |]
    in
    let idx = nonzero x in
    idx.(0), idx.(1)
  - : (int32, int32_elt) t * (int32, int32_elt) t =
  ([0, 1, 1, 2], [1, 0, 2, 2])

See also argwhere.

argwhere t is a 2-D tensor of shape [k; ndim t] whose rows are the coordinates of the k non-zero elements.

See also nonzero.

add a b is the element-wise sum of a and b.

add_s t s adds scalar s to each element of t.

radd_s s t is add_s t s.

sub a b is the element-wise difference a - b.

sub_s t s subtracts scalar s from each element.

rsub_s s t is s - t element-wise.

mul a b is the element-wise product of a and b.

mul_s t s multiplies each element by scalar s.

rmul_s s t is mul_s t s.

div a b is the element-wise quotient a / b.

Float dtypes use true division. Integer dtypes truncate toward zero.

  # let x = create int32 [| 2 |] [| -7l; 8l |] in
    let y = create int32 [| 2 |] [| 2l; 2l |] in
    div x y
  - : (int32, int32_elt) t = [-3, 4]

div_s t s divides each element by scalar s.

rdiv_s s t is s / t element-wise.

pow base exp is base raised to exp element-wise.

pow_s t s raises each element to scalar power s.

rpow_s s t is s{^t} element-wise.

mod_ a b is the element-wise remainder of a / b.

mod_s t s is the remainder of each element divided by scalar s.

rmod_s s t is s mod t element-wise.

neg t is the element-wise negation of t.

conjugate t is the complex conjugate of t. For complex dtypes, negates the imaginary part. For real dtypes, returns t unchanged.

abs t is the element-wise absolute value.

sign t is -1, 0, or 1 according to the sign of each element. For unsigned types, returns 1 for non-zero, 0 for zero.

  # create float32 [| 3 |] [| -2.; 0.; 3.5 |]
    |> sign
  - : (float, float32_elt) t = [-1, 0, 1]

square t is the element-wise square.

sqrt t is the element-wise square root.

rsqrt t is the element-wise reciprocal square root (1 / sqrt t).

recip t is the element-wise reciprocal (1 / t).

log t is the element-wise natural logarithm.

log2 t is the element-wise base-2 logarithm.

exp t is the element-wise exponential.

exp2 t is 2{^t} element-wise.

sin t is the element-wise sine.

cos t is the element-wise cosine.

tan t is the element-wise tangent.

asin t is the element-wise arcsine.

acos t is the element-wise arccosine.

atan t is the element-wise arctangent.

atan2 y x is the element-wise two-argument arctangent, returning angles in [-π, π].

sinh t is the element-wise hyperbolic sine.

cosh t is the element-wise hyperbolic cosine.

tanh t is the element-wise hyperbolic tangent.

asinh t is the element-wise inverse hyperbolic sine.

acosh t is the element-wise inverse hyperbolic cosine.

atanh t is the element-wise inverse hyperbolic tangent.

trunc t rounds each element toward zero.

ceil t rounds each element toward positive infinity.

floor t rounds each element toward negative infinity.

round t rounds each element to the nearest integer. Ties round away from zero (not banker's rounding).

  # create float32 [| 4 |] [| 2.5; 3.5; -2.5; -3.5 |]
    |> round
  - : (float, float32_elt) t = [3, 4, -3, -4]

hypot x y is sqrt(x² + y²) computed without intermediate overflow.

  # hypot (scalar float32 3.) (scalar float32 4.)
    |> item []
  - : float = 5.

lerp a b w is the linear interpolation a + w * (b - a). w is typically in [0, 1].

  # let a = create float32 [| 2 |] [| 1.; 2. |] in
    let b = create float32 [| 2 |] [| 5.; 8. |] in
    lerp a b (scalar float32 0.25)
  - : (float, float32_elt) t = [2, 3.5]

lerp_scalar_weight a b w is like lerp with a scalar weight.

isinf t is true where t is positive or negative infinity, false elsewhere. Non-float dtypes always return all false.

  # create float32 [| 4 |]
      [| 1.; Float.infinity;
         Float.neg_infinity; Float.nan |]
    |> isinf
  - : (bool, bool_elt) t = [false, true, true, false]

See also isnan, isfinite.

isnan t is true where t is NaN, false elsewhere. Non-float dtypes always return all false.

See also isinf, isfinite.

isfinite t is true where t is neither infinite nor NaN. Non-float dtypes always return all true.

See also isinf, isnan.

cmplt a b is true where a < b, false elsewhere.

less a b is cmplt.

less_s t s is true where t < s.

cmpne a b is true where a ≠ b, false elsewhere.

not_equal a b is cmpne.

not_equal_s t s is true where t ≠ s.

cmpeq a b is true where a = b, false elsewhere.

equal a b is cmpeq.

equal_s t s is true where t = s.

cmpgt a b is true where a > b, false elsewhere.

greater a b is cmpgt.

greater_s t s is true where t > s.

cmple a b is true where a ≤ b, false elsewhere.

less_equal a b is cmple.

less_equal_s t s is true where t ≤ s.

cmpge a b is true where a ≥ b, false elsewhere.

greater_equal a b is cmpge.

greater_equal_s t s is true where t ≥ s.

array_equal a b is a scalar true iff all elements of a and b are equal. Returns false if shapes differ.

  # let a = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    let b = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    array_equal a b |> item []
  - : bool = true

maximum a b is the element-wise maximum of a and b.

maximum_s t s is the element-wise maximum of t and scalar s.

rmaximum_s s t is maximum_s t s.

minimum a b is the element-wise minimum of a and b.

minimum_s t s is the element-wise minimum of t and scalar s.

rminimum_s s t is minimum_s t s.

logical_and a b is the element-wise logical AND. Non-zero is true.

logical_or a b is the element-wise logical OR.

logical_xor a b is the element-wise logical XOR.

logical_not t is the element-wise logical NOT: non-zero becomes 0, zero becomes 1.

where cond if_true if_false selects elements from if_true where cond is true and from if_false elsewhere. All three inputs broadcast to a common shape.

  # let x =
      create float32 [| 4 |] [| -1.; 2.; -3.; 4. |]
    in
    where
      (cmpgt x (scalar float32 0.))
      x (scalar float32 0.)
  - : (float, float32_elt) t = [0, 2, 0, 4]

clamp ?min ?max t clamps elements to [min, max]. Either bound may be omitted.

See also clip.

clip ?min ?max t is clamp.

bitwise_xor a b is the element-wise bitwise XOR.

bitwise_or a b is the element-wise bitwise OR.

bitwise_and a b is the element-wise bitwise AND.

bitwise_not t is the element-wise bitwise NOT.

invert t is bitwise_not.

lshift t n left-shifts each element by n bits.

Raises Invalid_argument if n is negative or the dtype is not an integer type.

  # create int32 [| 3 |] [| 1l; 2l; 3l |]
    |> Fun.flip lshift 2
  - : (int32, int32_elt) t = [4, 8, 12]

See also rshift.

rshift t n right-shifts each element by n bits.

Raises Invalid_argument if n is negative or the dtype is not an integer type.

See also lshift.

sum ?axes ?keepdims t sums elements along axes. When axes is omitted, reduces all axes (returns a scalar). When keepdims is true, reduced axes are kept with size 1. keepdims defaults to false. Negative axes count from the end.

  # create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |]
    |> sum |> item []
  - : float = 10.
  # create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |]
    |> sum ~axes:[ 0 ]
  - : (float, float32_elt) t = [4, 6]
  # create float32 [| 1; 2 |] [| 1.; 2. |]
    |> sum ~axes:[ 1 ] ~keepdims:true
  - : (float, float32_elt) t = float32 [1; 1] [[3]]

max ?axes ?keepdims t is the maximum along axes. NaN propagates. keepdims defaults to false.

  # create float32 [| 2; 3 |]
      [| 1.; 2.; 3.; 4.; 5.; 6. |]
    |> max |> item []
  - : float = 6.

min ?axes ?keepdims t is the minimum along axes. NaN propagates. keepdims defaults to false.

prod ?axes ?keepdims t is the product along axes. keepdims defaults to false.

  # create int32 [| 3 |] [| 2l; 3l; 4l |]
    |> prod |> item []
  - : int32 = 24l

cumsum ?axis t is the inclusive cumulative sum along axis. When axis is omitted, operates on the flattened tensor.

See also cumprod.

cumprod ?axis t is the inclusive cumulative product along axis. When axis is omitted, operates on the flattened tensor.

See also cumsum.

cummax ?axis t is the inclusive cumulative maximum along axis. NaN propagates for floating-point dtypes. When axis is omitted, operates on the flattened tensor.

See also cummin.

cummin ?axis t is the inclusive cumulative minimum along axis. NaN propagates for floating-point dtypes. When axis is omitted, operates on the flattened tensor.

See also cummax.

mean ?axes ?keepdims t is the arithmetic mean along axes. NaN propagates. keepdims defaults to false.

  # create float32 [| 4 |] [| 1.; 2.; 3.; 4. |]
    |> mean |> item []
  - : float = 2.5

var ?axes ?keepdims ?ddof t is the variance along axes. ddof (delta degrees of freedom) defaults to 0 (population variance); use 1 for sample variance. Computed as E[(X - E[X])²] / (N - ddof). keepdims defaults to false.

Raises Invalid_argument if ddof >= N.

  # create float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |]
    |> var |> item []
  - : float = 2.
  # create float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |]
    |> var ~ddof:1 |> item []
  - : float = 2.5

See also std.

std ?axes ?keepdims ?ddof t is the standard deviation: sqrt({!var} ~ddof t). ddof defaults to 0. keepdims defaults to false.

See also var.

all ?axes ?keepdims t is true iff every element along axes is non-zero. keepdims defaults to false.

  # create int32 [| 3 |] [| 1l; 2l; 3l |]
    |> all |> item []
  - : bool = true
  # create int32 [| 3 |] [| 1l; 0l; 3l |]
    |> all |> item []
  - : bool = false

See also any.

any ?axes ?keepdims t is true iff at least one element along axes is non-zero. keepdims defaults to false.

See also all.

argmax ?axis ?keepdims t is the index of the maximum along axis. Returns the first occurrence for ties. When axis is omitted, operates on the flattened tensor. keepdims defaults to false.

Raises Invalid_argument if axis is out of bounds.

  # create int32 [| 5 |] [| 3l; 1l; 4l; 1l; 5l |]
    |> argmax |> item []
  - : int32 = 4l

See also argmin.

argmin ?axis ?keepdims t is the index of the minimum along axis. Returns the first occurrence for ties. When axis is omitted, operates on the flattened tensor. keepdims defaults to false.

Raises Invalid_argument if axis is out of bounds.

See also argmax.

sort ?descending ?axis t sorts elements along axis and returns (sorted, indices) where indices maps sorted positions back to originals. descending defaults to false. axis defaults to -1 (last).

The sort is stable (equal elements preserve their relative order). NaN sorts to the end in ascending order and to the beginning in descending order.

Raises Invalid_argument if axis is out of bounds.

  # create int32 [| 5 |] [| 3l; 1l; 4l; 1l; 5l |]
    |> sort
  - : (int32, int32_elt) t * (int32, int32_elt) t =
  (int32 [5] [1, 1, ..., 4, 5], int32 [5] [1, 3, ..., 2, 4])

See also argsort.

argsort ?descending ?axis t is snd (sort ?descending ?axis t).

See also sort.

dot a b is the generalised dot product.

Contracts the last axis of a with:

• the only axis of b when b is 1-D,
• the second-to-last axis of b otherwise.

Dimension rules:

• 1-D × 1-D → scalar (inner product).
• 2-D × 2-D → matrix multiplication.
• N-D × M-D → contraction; output axes are the non-contracted axes of a followed by those of b.

Note. Unlike matmul, dot does not broadcast batch dimensions—it concatenates them.

Raises Invalid_argument if contraction axes differ in size or either input is 0-D.

  # let a = create float32 [| 2 |] [| 1.; 2. |] in
    let b = create float32 [| 2 |] [| 3.; 4. |] in
    dot a b |> item []
  - : float = 11.
  # dot (ones float32 [| 3; 4; 5 |])
        (ones float32 [| 5; 6 |]) |> shape
  - : int array = [|3; 4; 6|]

See also matmul, vdot, vecdot.

matmul a b is the matrix product of a and b with batch broadcasting.

Dimension rules:

• 1-D × 1-D → scalar (inner product).
• 1-D × N-D → a is treated as a row vector.
• N-D × 1-D → b is treated as a column vector.
• N-D × M-D → matrix multiply on last two axes; leading axes are broadcast.

Raises Invalid_argument if inputs are 0-D or inner dimensions mismatch.

  # let a =
      create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |]
    in
    let b = create float32 [| 2 |] [| 5.; 6. |] in
    matmul a b
  - : (float, float32_elt) t = [17, 39]
  # matmul (ones float32 [| 1; 3; 4 |])
           (ones float32 [| 5; 4; 2 |]) |> shape
  - : int array = [|5; 3; 2|]

See also dot, multi_dot.

diagonal ?offset ?axis1 ?axis2 t extracts diagonals from 2-D planes defined by axis1 and axis2. offset defaults to 0. axis1 and axis2 default to the last two axes.

Raises Invalid_argument if axis1 = axis2 or either is out of bounds.

See also diag, trace.

matrix_transpose t swaps the last two axes: […; m; n] → […; n; m]. For 1-D tensors, returns t unchanged.

See also transpose.

vdot a b is the dot product of two vectors. Both inputs are flattened; for complex dtypes, a is conjugated first. Always returns a scalar.

Raises Invalid_argument if the inputs have different numbers of elements.

See also dot, vecdot.

vecdot ?axis a b is the dot product of a and b along axis with broadcasting. axis defaults to -1.

Raises Invalid_argument if the specified axis dimensions differ.

See also vdot, dot.

inner a b is the inner product over the last axes of a and b.

Raises Invalid_argument if the last dimensions differ.

See also dot, outer.

outer a b is the outer product. Inputs are flattened to 1-D; the result has shape [numel a; numel b].

See also inner.

tensordot ?axes a b contracts a and b along the specified axis pairs. axes defaults to contracting the last axis of a with the first axis of b.

Raises Invalid_argument if the contracted axes have different sizes.

einsum subscripts operands evaluates Einstein summation.

  # let a =
      create float32 [| 2; 3 |]
        [| 1.; 2.; 3.; 4.; 5.; 6. |]
    in
    let b =
      create float32 [| 3; 2 |]
        [| 1.; 2.; 3.; 4.; 5.; 6. |]
    in
    einsum "ij,jk->ik" [| a; b |] |> shape
  - : int array = [|2; 2|]

See also matmul, tensordot.

kron a b is the Kronecker product. The result has shape [a.shape.(i) * b.shape.(i)] for each i.

multi_dot ts is the chained matrix product of ts, automatically choosing the association order that minimises computation.

Raises Invalid_argument if the array is empty, shapes are incompatible, or dtypes are not floating-point or complex.

See also matmul.

matrix_power t n raises square matrix t to integer power n. n = 0 returns the identity; n < 0 uses the inverse.

Raises Invalid_argument if t is not square, the dtype is not floating-point or complex, or n < 0 and t is singular.

cross ?axis a b is the cross product of 3-element vectors along axis. axis defaults to -1.

Raises Invalid_argument if the axis dimension is not 3.

cholesky ?upper a is the Cholesky factor of positive- definite matrix a. When upper is true, returns the upper-triangular factor U such that a = Uᵀ U; otherwise (default) returns the lower-triangular factor L such that a = L Lᵀ.

Raises Invalid_argument if a is not positive-definite or the dtype is not floating-point or complex.

See also solve.

qr ?mode a is (Q, R) where a = Q R, Q is orthogonal, and R is upper-triangular. mode defaults to `Reduced.

Raises Invalid_argument if the dtype is not floating-point or complex.

See also svd.