Module RuneSource

('a, 'b) t is a tensor with OCaml type 'a and bigarray type 'b.

Index specification for tensor slicing

data t returns underlying bigarray buffer.

Buffer may contain data beyond tensor bounds for strided views. Direct access requires careful index computation using strides and offset.

• raises Invalid_argument

  if tensor not C-contiguous when safety matters

shape t returns dimensions. Empty array for scalars.

dtype t returns data type.

strides t returns byte strides for each dimension.

stride i t returns byte stride for dimension i.

• raises Invalid_argument

  if i out of bounds

dims t is synonym for shape.

dim i t returns size of dimension i.

• raises Invalid_argument

  if i out of bounds

ndim t returns number of dimensions.

itemsize t returns bytes per element.

size t returns total number of elements.

numel t is synonym for size.

nbytes t returns size t * itemsize t.

offset t returns element offset in underlying buffer.

is_c_contiguous t returns true if elements are contiguous in C order.

to_bigarray t converts to standard bigarray.

Always returns contiguous copy with same shape. Use for interop with libraries expecting standard bigarrays.

  # let t = create float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |]
  val t : (float, float32_elt) t = [[1, 2, 3],
                                    [4, 5, 6]]
  # shape (to_bigarray t |> of_bigarray)
  - : int array = [|2; 3|]

• raises Failure

  if tensor dtype is an extended type not supported by standard Bigarray

to_bigarray_ext t converts to extended bigarray.

Always returns contiguous copy with same shape. Use for interop with libraries expecting extended bigarrays (e.g., with bfloat16 support).

  # let t = create float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |]
  val t : (float, float32_elt) t = [[1, 2, 3],
                                    [4, 5, 6]]
  # shape (to_bigarray_ext t |> of_bigarray_ext)
  - : int array = [|2; 3|]

to_array t converts to OCaml array.

Flattens tensor to 1-D array in row-major (C) order. Always copies.

  # let t = create int32 [| 2; 2 |] [| 1l; 2l; 3l; 4l |]
  val t : (int32, int32_elt) t = [[1, 2],
                                  [3, 4]]
  # to_array t
  - : int32 array = [|1l; 2l; 3l; 4l|]

create dtype shape data creates tensor from array data.

Length of data must equal product of shape.

• raises Invalid_argument

  if array size doesn't match shape

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

init dtype shape f creates tensor where element at indices i has value f i.

Function f receives array of indices for each position. Useful for creating position-dependent values.

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

  # init float32 [| 3; 3 |] (fun i -> if i.(0) = i.(1) then 1. else 0.)
  - : (float, float32_elt) t = [[1, 0, 0],
                                [0, 1, 0],
                                [0, 0, 1]]

empty dtype shape allocates uninitialized tensor.

full dtype shape value creates tensor filled with value.

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

ones dtype shape creates tensor filled with ones.

zeros dtype shape creates tensor filled with zeros.

scalar dtype value creates scalar tensor containing value.

empty_like t creates uninitialized tensor with same shape and dtype as t.

full_like t value creates tensor shaped like t filled with value.

ones_like t creates tensor shaped like t filled with ones.

zeros_like t creates tensor shaped like t filled with zeros.

scalar_like t value creates scalar with same dtype as t.

eye ?m ?k dtype n creates matrix with ones on k-th diagonal.

Default m = n (square), k = 0 (main diagonal). Positive k shifts diagonal above main, negative below.

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

identity dtype n creates n×n identity matrix.

Equivalent to eye dtype n. Square matrix with ones on main diagonal, zeros elsewhere.

  # identity int32 3
  - : (int32, int32_elt) t = [[1, 0, 0],
                              [0, 1, 0],
                              [0, 0, 1]]

diag ?k v extracts diagonal or constructs diagonal array.

If v is 1D, returns 2D array with v on the k-th diagonal. If v is 2D, returns 1D array containing the k-th diagonal. Use k > 0 for diagonals above the main diagonal, k < 0 for diagonals below.

• parameter k

  Diagonal offset (default 0 = main diagonal)

• raises Failure

  if v is 0D

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

arange dtype start stop step generates values from start to [stop).

Step must be non-zero. Result length is (stop - start) / step rounded toward zero.

• raises Failure

  if step = 0

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

arange_f dtype start stop step generates float values from start to [stop).

Like arange but for floating-point ranges. Handles fractional steps. Due to floating-point precision, final value may differ slightly from expected.

• raises Failure

  if step = 0.0

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

linspace dtype ?endpoint start stop count generates count evenly spaced values from start to stop.

If endpoint is true (default), stop is included.

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

logspace dtype ?endpoint ?base start_exp stop_exp count generates values evenly spaced on log scale.

Returns base ** x where x ranges from start_exp to stop_exp. Default base = 10.0.

  # 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]

geomspace dtype ?endpoint start stop count generates values evenly spaced on geometric (multiplicative) scale.

• raises Invalid_argument

  if start <= 0. or stop <= 0.

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

meshgrid ?indexing x y creates coordinate grids from 1D arrays.

Returns (X, Y) where X and Y are 2D arrays representing grid coordinates.

• `xy (default): Cartesian indexing - X changes along columns, Y changes along rows
• `ij: Matrix indexing - X changes along rows, Y changes along columns

• raises Invalid_argument

  if x or y are not 1D

  # 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 =
  ([[0, 1, 2],
    [0, 1, 2]], [[0, 0, 0],
                 [1, 1, 1]])

tril ?k x returns lower triangular part of matrix.

Elements above the k-th diagonal are zeroed.

• k = 0 (default): main diagonal
• k > 0: include k diagonals above main
• k < 0: exclude |k| diagonals below main

• raises Invalid_argument

  if x has less than 2 dimensions

triu ?k x returns upper triangular part of matrix.

Elements below the k-th diagonal are zeroed.

• k = 0 (default): main diagonal
• k > 0: exclude k diagonals above main
• k < 0: include |k| diagonals below main

• raises Invalid_argument

  if x has less than 2 dimensions

of_bigarray ba creates tensor from standard bigarray.

Zero-copy when bigarray is contiguous. Creates view sharing same memory. Modifications to either affect both.

  # let t = zeros float32 [| 2; 3 |] in
    t
  - : (float, float32_elt) t = [[0, 0, 0],
                                [0, 0, 0]]

of_bigarray_ext ba creates tensor from extended bigarray.

Zero-copy when bigarray is contiguous. Creates view sharing same memory. Modifications to either affect both. Supports extended types like bfloat16.

  # let ba = Bigarray_ext.Genarray.create Bigarray_ext.bfloat16
              Bigarray_ext.c_layout [| 2; 3 |] in
    let t = of_bigarray_ext ba in
    shape t
  - : int array = [|2; 3|]

rand dtype ?seed shape generates uniform random values in [0, 1).

Only supports float dtypes. Same seed produces same sequence.

• raises Invalid_argument

  if non-float dtype

randn dtype ?seed shape generates standard normal random values.

Mean 0, variance 1. Uses Box-Muller transform for efficiency. Only supports float dtypes. Same seed produces same sequence.

• raises Invalid_argument

  if non-float dtype

randint dtype ?seed ?high shape low generates integers in [low, high).

Uniform distribution over range. Default high = 10. Note: high is exclusive (NumPy convention).

• raises Invalid_argument

  if non-integer dtype or low >= high

dropout ?seed ~rate x randomly zeroes elements with probability rate.

The remaining values are scaled by 1/(1 - rate) keeping the expected value unchanged. Requires floating-point input and rate in [0, 1). When seed is provided, the same mask is reproduced.

reshape shape t returns view with new shape.

At most one dimension can be -1 (inferred from total elements). Product of dimensions must match total elements. Returns a zero-copy view when the new layout is compatible; raises otherwise.

• raises Invalid_argument

  if shape incompatible or multiple -1 dimensions

• raises Invalid_argument

  if tensor layout cannot support requested reshape

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

broadcast_to shape t broadcasts tensor to target shape.

Shapes must be broadcast-compatible: dimensions align from right, each must be equal or source must be 1. Returns view (no copy) with zero strides for broadcast dimensions.

• raises Invalid_argument

  if shapes incompatible

  # let t = create int32 [|1; 3|] [|1l; 2l; 3l|] in
    broadcast_to [|3; 3|] t
  - : (int32, int32_elt) t = [[1, 2, 3],
                              [1, 2, 3],
                              [1, 2, 3]]
  # let t = ones float32 [|3; 1|] in
    shape (broadcast_to [|2; 3; 4|] t)
  - : int array = [|2; 3; 4|]

broadcasted ?reverse t1 t2 broadcasts tensors to common shape.

Returns views of both tensors broadcast to compatible shape. If reverse is true, returns (t2', t1') instead of (t1', t2'). Useful before element-wise operations.

• raises Invalid_argument

  if shapes incompatible

  # let t1 = ones float32 [|3;1|] in
    let t2 = ones float32 [|1;5|] in
    let t1', t2' = broadcasted t1 t2 in
    shape t1', shape t2'
  - : int array * int array = ([|3; 5|], [|3; 5|])

expand shape t broadcasts tensor where -1 keeps original dimension.

Like broadcast_to but -1 preserves existing dimension size. Adds dimensions on left if needed.

  # let t = ones float32 [|1; 4; 1|] in
    shape (expand [|3; -1; 5|] t)
  - : int array = [|3; 4; 5|]
  # let t = ones float32 [|5; 5|] in
    shape (expand [|-1; -1|] t)
  - : int array = [|5; 5|]

flatten ?start_dim ?end_dim t collapses dimensions into single dimension.

Default start_dim = 0, end_dim = -1 (last). Negative indices count from end. Dimensions start_dim through end_dim inclusive are flattened.

• raises Invalid_argument

  if indices out of bounds

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

unflatten dim sizes t expands dimension dim into multiple dimensions.

Product of sizes must equal size of dimension dim. At most one dimension can be -1 (inferred). Inverse of flatten.

• raises Invalid_argument

  if product mismatch or dim out of bounds

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

ravel t returns contiguous 1-D view.

Equivalent to flatten t but always returns contiguous result. Use when you need both flattening and contiguity.

  # let x = create int32 [| 2; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l |] in
    ravel x
  - : (int32, int32_elt) t = [1, 2, 3, 4, 5, 6]
  # let t = transpose (ones float32 [| 3; 4 |]) in
    is_c_contiguous t
  - : bool = false
  # let t_ravel = ravel t in
    is_c_contiguous t_ravel
  - : bool = true

squeeze ?axes t removes dimensions of size 1.

If axes specified, only removes those dimensions. Negative indices count from end. Returns view when possible.

• raises Invalid_argument

  if specified axis doesn't have size 1

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

unsqueeze ?axes t inserts dimensions of size 1 at specified positions.

Axes refer to positions in result tensor. Must be in range 0, ndim.

• raises Invalid_argument

  if axes not specified, out of bounds, or contains duplicates

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

squeeze_axis axis t removes dimension axis if size is 1.

• raises Invalid_argument

  if dimension size is not 1

unsqueeze_axis axis t inserts dimension of size 1 at axis.

expand_dims axes t is synonym for unsqueeze.

transpose ?axes t permutes dimensions.

Default reverses all dimensions. axes must be permutation of 0..ndim-1. Returns view (no copy) with adjusted strides.

• raises Invalid_argument

  if axes not valid permutation

  # let x = create int32 [| 2; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l |] in
    transpose x
  - : (int32, int32_elt) t = [[1, 4],
                              [2, 5],
                              [3, 6]]
  # transpose ~axes:[ 2; 0; 1 ] (zeros float32 [| 2; 3; 4 |]) |> shape
  - : int array = [|4; 2; 3|]
  # let id = transpose ~axes:[ 1; 0 ] in
    id == transpose
  - : bool = false

flip ?axes t reverses order along specified dimensions.

Default flips all dimensions.

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

as_strided shape strides ~offset t creates a strided view of the tensor with custom shape, strides (in elements), and offset (in elements).

This is a low-level operation that allows arbitrary memory layouts. Backends that support arbitrary strides (e.g., native) implement this as zero-copy. Other backends may need to materialize the data.

Warning: This function can create views that access out-of-bounds memory if used incorrectly. Use with caution.

  # let x = create float32 [| 8 |] [| 0.; 1.; 2.; 3.; 4.; 5.; 6.; 7. |] in
    as_strided [| 3; 3 |] [| 2; 1 |] ~offset:0 x
  - : (float, float32_elt) t = [[0, 1, 2],
                                [2, 3, 4],
                                [4, 5, 6]]

moveaxis src dst t moves dimension from src to dst.

• raises Invalid_argument

  if indices out of bounds

swapaxes axis1 axis2 t exchanges two dimensions.

• raises Invalid_argument

  if indices out of bounds

roll ?axis shift t shifts elements along axis.

Elements shifted beyond last position wrap to beginning. If axis not specified, shifts flattened tensor. Negative shift rolls backward.

• raises Invalid_argument

  if axis out of bounds

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

pad padding value t pads tensor with value.

padding specifies (before, after) for each dimension. Length must match tensor dimensions. Negative padding not allowed.

• raises Invalid_argument

  if padding length wrong or negative values

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

shrink ranges t extracts slice from start to stop (exclusive) for each dimension.

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

tile reps t constructs tensor by repeating t.

reps specifies repetitions per dimension. If longer than ndim, prepends dimensions. Zero repetitions create empty tensor.

• raises Invalid_argument

  if reps contains negative values

  # let x = create int32 [| 1; 2 |] [| 1l; 2l |] in
    tile [| 2; 3 |] x
  - : (int32, int32_elt) t = [[1, 2, 1, 2, 1, 2],
                              [1, 2, 1, 2, 1, 2]]
  # let x = create int32 [| 2 |] [| 1l; 2l |] in
    tile [| 2; 1; 3 |] x |> shape
  - : int array = [|2; 1; 6|]

repeat ?axis count t repeats elements count times.

If axis not specified, repeats flattened tensor.

  # let x = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    repeat 2 x
  - : (int32, int32_elt) t = [1, 1, 2, 2, 3, 3]
  # let x = create int32 [| 1; 2 |] [| 1l; 2l |] in
    repeat ~axis:0 3 x
  - : (int32, int32_elt) t = [[1, 2],
                              [1, 2],
                              [1, 2]]

concatenate ?axis ts joins tensors along existing axis.

All tensors must have same shape except on concatenation axis. If axis not specified, flattens all tensors then concatenates. Returns contiguous result.

• raises Invalid_argument

  if empty list or shape mismatch

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

stack ?axis ts joins tensors along new axis.

All tensors must have identical shape. Result rank is input rank + 1. Default axis=0. Negative axis counts from end of result shape.

• raises Invalid_argument

  if empty list, shape mismatch, or axis out of bounds

  # let x1 = create int32 [| 2 |] [| 1l; 2l |] in
    let x2 = create int32 [| 2 |] [| 3l; 4l |] in
    stack [x1; x2]
  - : (int32, int32_elt) t = [[1, 2],
                              [3, 4]]
  # let x1 = create int32 [| 2 |] [| 1l; 2l |] in
    let x2 = create int32 [| 2 |] [| 3l; 4l |] in
    stack ~axis:1 [x1; x2]
  - : (int32, int32_elt) t = [[1, 3],
                              [2, 4]]
  # stack ~axis:(-1) [ones float32 [| 2; 3 |]; zeros float32 [| 2; 3 |]] |> shape
  - : int array = [|2; 3; 2|]

vstack ts stacks tensors vertically (row-wise).

1-D tensors are treated as row vectors (shape 1;n). Higher-D tensors concatenate along axis 0. All tensors must have same shape except possibly first dimension.

• raises Invalid_argument

  if incompatible shapes

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

hstack ts stacks tensors horizontally (column-wise).

1-D tensors concatenate directly. Higher-D tensors concatenate along axis 1. For 1-D arrays of different lengths, use vstack to make 2-D first.

• raises Invalid_argument

  if incompatible shapes or <2D with different axis 0

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

dstack ts stacks tensors depth-wise (along third axis).

Tensors are reshaped to at least 3-D before concatenation:

• 1-D shape n → 1;n;1
• 2-D shape m;n → m;n;1
• 3-D+ unchanged

• raises Invalid_argument

  if resulting shapes incompatible

  # let x1 = create int32 [| 2 |] [| 1l; 2l |] in
    let x2 = create int32 [| 2 |] [| 3l; 4l |] in
    dstack [x1; x2]
  - : (int32, int32_elt) t = [[[1, 3],
                               [2, 4]]]
  # let x1 = create int32 [| 2; 2 |] [| 1l; 2l; 3l; 4l |] in
    let x2 = create int32 [| 2; 2 |] [| 5l; 6l; 7l; 8l |] in
    dstack [x1; x2]
  - : (int32, int32_elt) t = [[[1, 5],
                               [2, 6]],
                              [[3, 7],
                               [4, 8]]]

broadcast_arrays ts broadcasts all tensors to common shape.

Finds the common broadcast shape and returns list of views with that shape. Broadcasting rules: dimensions align right, each must be 1 or equal.

• raises Invalid_argument

  if shapes incompatible or empty list

  # let x1 = ones float32 [| 3; 1 |] in
    let x2 = ones float32 [| 1; 5 |] in
    broadcast_arrays [x1; x2] |> List.map shape
  - : int array list = [[|3; 5|]; [|3; 5|]]
  # let x1 = scalar float32 5. in
    let x2 = ones float32 [| 2; 3; 4 |] in
    broadcast_arrays [x1; x2] |> List.map shape
  - : int array list = [[|2; 3; 4|]; [|2; 3; 4|]]

array_split ~axis sections t splits tensor into multiple parts.

`Count n divides into n parts as evenly as possible. Extra elements go to first parts. `Indices [i1;i2;...] splits at indices creating start:i1, i1:i2, i2:end.

• raises Invalid_argument

  if axis out of bounds or invalid sections

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

split ~axis sections t splits into equal parts.

• raises Invalid_argument

  if axis size not divisible by sections

  # let x = create int32 [| 4; 2 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l |] in
    split ~axis:0 2 x
  - : (int32, int32_elt) t list = [[[1, 2],
                                    [3, 4]]; [[5, 6],
                                              [7, 8]]]

cast dtype t converts elements to new dtype.

Returns copy with same values in new type.

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

astype dtype t is synonym for cast.

contiguous t returns C-contiguous tensor.

Returns t unchanged if already contiguous (O(1)), otherwise creates contiguous copy (O(n)). Use before operations requiring direct memory access.

  # let t = transpose (ones float32 [| 3; 4 |]) in
    is_c_contiguous (contiguous t)
  - : bool = true

copy t returns deep copy.

Always allocates new memory and copies data. 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])

blit src dst copies src into dst.

Shapes must match exactly. Handles broadcasting internally. Modifies dst in-place.

• raises Invalid_argument

  if shape mismatch

  let dst = zeros float32 [| 3; 3 |] in
  blit (ones float32 [| 3; 3 |]) dst
  (* dst now contains all 1s *)

ifill value t sets all elements of t to value in-place.

fill value t returns a copy of t filled with value, leaving t unchanged. Handy when wanting a filled tensor without mutating the source.

get indices t returns subtensor at indices.

Indexes from outermost dimension. Returns scalar tensor if all dimensions indexed, otherwise returns view of remaining dimensions.

• raises Invalid_argument

  if indices out of bounds

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

set indices t value assigns value at indices.

• raises Invalid_argument

  if indices out of bounds

slice specs t extracts subtensor using advanced indexing.

Each element in specs corresponds to an axis from left to right:

• I i: single index (reduces dimension; negative from end)
• L [i1;i2;...]: fancy indexing with list of indices
• R (start, stop): range [start, stop) with step 1
• Rs (start, stop, step): range with step
• A: full axis (default for missing specs)
• M mask: boolean mask (shape must match axis)
• N: insert new dimension of size 1

Missing specs default to A. Returns view when possible.

• raises Invalid_argument

  if specs out of bounds or incompatible

  # let x = create int32 [| 2; 4 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l |] in
    slice [ I 1 ] x
  - : (int32, int32_elt) t = [5, 6, 7, 8]
  # let x = create int32 [| 5 |] [| 0l; 1l; 2l; 3l; 4l |] in
    slice [ R (1, 3) ] x
  - : (int32, int32_elt) t = [1, 2]
  # 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 = [[1, 3],
                              [4, 6]]
  # let x = create int32 [| 3; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |] in
    slice [ A; N ] x  (* Add new axis at position 1 *)
  - : (int32, int32_elt) t = [[[1, 2, 3]],
                              [[4, 5, 6]],
                              [[7, 8, 9]]]

set_slice specs t value assigns value to indexed region.

Like slice but modifies t in-place. Value is broadcast if needed.

• raises Invalid_argument

  if specs incompatible

item indices t returns scalar value at indices.

Must provide indices for all dimensions.

• raises Invalid_argument

  if wrong number of indices or out of bounds

set_item indices value t sets scalar value at indices.

Must provide indices for all dimensions. Modifies tensor in-place.

• raises Invalid_argument

  if wrong number of indices or out of bounds

take ?axis ?mode indices t takes elements from t using indices.

Equivalent to tindices in NumPy along the specified axis. If axis is None, flattens t first. indices is an integer tensor of indices to take. mode handles out-of-bounds indices: `raise (default), `wrap (modulo), `clip (clamp to bounds).

Returns a new tensor with shape based on indices and t's shape.

• raises Invalid_argument

  if indices out of bounds and mode=`raise

  # 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]
  # let x = create int32 [| 3; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |] in
    take ~axis:1 (create int32 [| 3 |] [| 0l; 2l; 1l |]) x
  - : (int32, int32_elt) t = [[1, 3, 2],
                              [4, 6, 5],
                              [7, 9, 8]]
  # let x = create int32 [| 5 |] [| 0l; 1l; 2l; 3l; 4l |] in
    take ~mode:`clip (create int32 [| 2 |] [| -1l; 5l |]) x  (* Clamps to [0,4] *)
  - : (int32, int32_elt) t = [0, 4]

take_along_axis ~axis indices t takes values along the specified axis using indices.

Equivalent to NumPy's take_along_axis. indices must have the same shape as t except along the specified axis, where it matches the output size. Useful for gathering from argmax/argmin results.

• raises Invalid_argument

  if shapes incompatible

  # let x = create float32 [| 2; 3 |] [| 4.; 1.; 2.; 3.; 5.; 6. |] in
    let indices = create int32 [| 2; 1 |] [| 1l; 0l |] in  (* Per row indices *)
    take_along_axis ~axis:1 indices x
  - : (float, float32_elt) t = [[1],
                                [3]]
  # let x = create float32 [| 3; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8.; 9. |] in
    let indices = expand_dims [ 0 ] (argmax ~axis:0 x) in  (* Shape [1, 3] *)
    take_along_axis ~axis:0 indices x  (* Max per column *)
  - : (float, float32_elt) t = [[7, 8, 9]]

put ?axis ~indices ~values ?mode t sets elements in t at positions specified by indices to values.

Equivalent to NumPy's put (in-place version of take for setting). If axis is None, flattens t first. indices is an integer tensor of positions to set. values must match the number of indices (broadcasted if needed). mode handles out-of-bounds indices: `raise (default), `wrap (modulo), `clip (clamp).

Modifies t in-place.

• raises Invalid_argument

  if shapes incompatible or indices out of bounds with mode=`raise

  # let x = zeros int32 [| 5 |] in
    put ~indices:(create int32 [| 3 |] [| 1l; 3l; 0l |])
        ~values:(create int32 [| 3 |] [| 10l; 20l; 30l |]) x;
    x
  - : (int32, int32_elt) t = [30, 10, 0, 20, 0]
  # let x = create int32 [| 3; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |] in
    put ~axis:1 ~indices:(create int32 [| 3; 1 |] [| 0l; 2l; 1l |])
                ~values:(create int32 [| 3; 1 |] [| 10l; 20l; 30l |]) x;
    x
  - : (int32, int32_elt) t = [[10, 2, 3],
                              [4, 5, 20],
                              [7, 30, 9]]
  # let y = zeros int32 [| 5 |] in
    put ~mode:`clip ~indices:(create int32 [| 2 |] [| -1l; 5l |])
        ~values:(create int32 [| 2 |] [| 99l; 99l |]) y;
    y  (* Clamps to [0,4] *)
  - : (int32, int32_elt) t = [99, 0, 0, 0, 99]

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

indices is an array that contains one tensor per axis of t. Each tensor provides integer coordinates for its axis; they are broadcast to a common shape that also determines how many updates are performed. values is broadcast to the same shape. Updates follow element-wise order and leave the shape of t unchanged. Duplicate coordinates overwrite previous updates, matching put.

mode controls how out-of-bounds indices are handled per axis: `raise (default) checks bounds, `wrap performs modular indexing, and `clip clamps to the valid range.

• raises Invalid_argument

  if the number of index tensors does not match the rank of t, or if any axis is zero-sized while a non-empty update set is requested.

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

put_along_axis ~axis ~indices ~values t sets values along the specified axis using indices.

Equivalent to NumPy's put_along_axis. indices must have the same shape as t except along the axis (where it matches values' size along that axis). values is broadcasted to match the selection shape. Useful for scattering to argmax/argmin positions.

Modifies t in-place.

• raises Invalid_argument

  if shapes incompatible

  # let x = zeros float32 [| 2; 3 |] in
    let indices = create int32 [| 2; 1 |] [| 1l; 0l |] in  (* Per row positions *)
    put_along_axis ~axis:1 ~indices ~values:(create float32 [| 2; 1 |] [| 10.; 20. |]) x;
    x
  - : (float, float32_elt) t = [[0, 10, 0],
                                [20, 0, 0]]
  # let x = create float32 [| 3; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8.; 9. |] in
    let indices = expand_dims [ 0 ] (argmax ~axis:0 x) in  (* Shape [1, 3] *)
    put_along_axis ~axis:0 ~indices ~values:(ones float32 [| 1; 3 |]) x;
    x  (* Set max per column to 1 *)
  - : (float, float32_elt) t = [[1, 2, 3],
                                [4, 5, 6],
                                [1, 1, 1]]

compress ?axis condition t selects elements where condition is true.

Equivalent to NumPy's compress. condition is a 1D boolean array. If axis is None, flattens t first. Otherwise, compresses along the specified axis (condition length must match t's dim along axis).

Returns a new tensor with reduced size along the axis/flattened.

• raises Invalid_argument

  if condition length 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]
  # let x = create int32 [| 3; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |] in
    compress ~axis:0 ~condition:(create bool [| 3 |] [| false; true; true |]) x
  - : (int32, int32_elt) t = [[4, 5, 6],
                              [7, 8, 9]]

extract condition t flattens and selects elements where condition is true.

Equivalent to NumPy's extract (1D compress after flatten). condition must have the same shape and size as t (element-wise).

Returns a 1D tensor with selected elements.

• raises Invalid_argument

  if shapes incompatible

  # let x = create int32 [| 3; 3 |] [| 1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l |] in
    extract ~condition:(greater_s x 5l) x
  - : (int32, int32_elt) t = [6, 7, 8, 9]

nonzero t returns indices of non-zero elements.

Equivalent to NumPy's nonzero. Treats non-zero as true for bool tensors. Returns an array of 1D tensors, one per dimension, with coordinates of non-zeros.

For example, for a 2D tensor, returns | rows; cols | where rowsi, colsi is the position of the i-th non-zero.

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

argwhere t returns indices of non-zero elements as a 2D tensor.

Equivalent to NumPy's argwhere. Each row is a coordinate dim0; dim1; ... of a non-zero element. Shape is num_nonzeros; ndim.

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

add t1 t2 computes element-wise sum with broadcasting.

• raises Invalid_argument

  if shapes incompatible

add_s t scalar adds scalar to each element.

iadd target value adds value to target in-place.

Returns modified target.

radd_s scalar t is add_s t scalar.

iadd_s t scalar adds scalar to t in-place.

sub t1 t2 computes element-wise difference with broadcasting.

sub_s t scalar subtracts scalar from each element.

rsub_s scalar t computes scalar - t.

isub target value subtracts value from target in-place.

isub_s t scalar subtracts scalar from t in-place.

mul t1 t2 computes element-wise product with broadcasting.

mul_s t scalar multiplies each element by scalar.

rmul_s scalar t is mul_s t scalar.

imul target value multiplies target by value in-place.

imul_s t scalar multiplies t by scalar in-place.

div t1 t2 computes element-wise division.

True division for floats (result is float). Integer division for integers (truncates toward zero). Complex division follows standard rules.

  # let x = create float32 [| 3 |] [| 7.; 8.; 9. |] in
    let y = create float32 [| 3 |] [| 2.; 2.; 2. |] in
    div x y
  - : (float, float32_elt) t = [3.5, 4, 4.5]
  # let x = create int32 [| 3 |] [| 7l; 8l; 9l |] in
    let y = create int32 [| 3 |] [| 2l; 2l; 2l |] in
    div x y
  - : (int32, int32_elt) t = [3, 4, 4]
  # 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 scalar divides each element by scalar.

rdiv_s scalar t computes scalar / t.

idiv target value divides target by value in-place.

idiv_s t scalar divides t by scalar in-place.

pow base exponent computes element-wise power.

pow_s t scalar raises each element to scalar power.

rpow_s scalar t computes scalar ** t.

ipow target exponent raises target to exponent in-place.

ipow_s t scalar raises t to scalar power in-place.

mod_ t1 t2 computes element-wise modulo.

mod_s t scalar computes modulo scalar for each element.

rmod_s scalar t computes scalar mod t.

imod target divisor computes modulo in-place.

imod_s t scalar computes modulo scalar in-place.

neg t negates all elements.

conjugate x computes the complex conjugate.

For complex tensors, negates the imaginary part of each element. For real tensors, returns the input unchanged.

  # let x = create complex32 [| 2 |]
      [|Complex.{re=1.; im=2.}; Complex.{re=3.; im=4.}|] in
    conjugate x |> to_array
  - : Complex.t array =
  [|{Complex.re = 1.; im = -2.}; {Complex.re = 3.; im = -4.}|]

abs t computes absolute value.

sign t returns -1, 0, or 1 based on sign.

For unsigned types, returns 1 for all non-zero values, 0 for zero.

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

square t computes element-wise square.

sqrt t computes element-wise square root.

rsqrt t computes reciprocal square root.

recip t computes element-wise reciprocal.

log t computes natural logarithm.

log2 t computes base-2 logarithm.

exp t computes exponential.

exp2 t computes 2^x.

sin t computes sine.

cos t computes cosine.

tan t computes tangent.

asin t computes arcsine.

acos t computes arccosine.

atan t computes arctangent.

atan2 y x computes arctangent of y/x using signs to determine quadrant.

Returns angle in radians in range -π, π. Handles x=0 correctly.

  # let y = scalar float32 1. in
    let x = scalar float32 1. in
    atan2 y x |> item [] |> Float.round
  - : float = 1.
  # let y = scalar float32 1. in
    let x = scalar float32 0. in
    atan2 y x |> item [] |> Float.round
  - : float = 2.
  # let y = scalar float32 0. in
    let x = scalar float32 0. in
    atan2 y x |> item []
  - : float = 0.

sinh t computes hyperbolic sine.

cosh t computes hyperbolic cosine.

tanh t computes hyperbolic tangent.

asinh t computes inverse hyperbolic sine.

acosh t computes inverse hyperbolic cosine.

atanh t computes inverse hyperbolic tangent.

hypot x y computes sqrt(x² + y²) avoiding overflow.

Uses numerically stable algorithm: max * sqrt(1 + (min/max)²).

  # let x = scalar float32 3. in
    let y = scalar float32 4. in
    hypot x y |> item []
  - : float = 5.
  # let x = scalar float64 1e200 in
    let y = scalar float64 1e200 in
    hypot x y |> item [] < Float.infinity
  - : bool = true

trunc t rounds toward zero.

Removes fractional part. Positive values round down, negative round up.

  # let x = create float32 [| 3 |] [| 2.7; -2.7; 2.0 |] in
    trunc x
  - : (float, float32_elt) t = [2, -2, 2]

ceil t rounds up to nearest integer.

Smallest integer not less than input.

  # let x = create float32 [| 4 |] [| 2.1; 2.9; -2.1; -2.9 |] in
    ceil x
  - : (float, float32_elt) t = [3, 3, -2, -2]

floor t rounds down to nearest integer.

Largest integer not greater than input.

  # let x = create float32 [| 4 |] [| 2.1; 2.9; -2.1; -2.9 |] in
    floor x
  - : (float, float32_elt) t = [2, 2, -3, -3]

round t rounds to nearest integer (half away from zero).

Ties round away from zero (not banker's rounding).

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

lerp start end_ weight computes linear interpolation.

Returns start + weight * (end_ - start). Weight typically in 0, 1.

  # let start = scalar float32 0. in
    let end_ = scalar float32 10. in
    let weight = scalar float32 0.3 in
    lerp start end_ weight |> item []
  - : float = 3.
  # let start = create float32 [| 2 |] [| 1.; 2. |] in
    let end_ = create float32 [| 2 |] [| 5.; 8. |] in
    let weight = create float32 [| 2 |] [| 0.25; 0.5 |] in
    lerp start end_ weight
  - : (float, float32_elt) t = [2, 5]

lerp_scalar_weight start end_ weight interpolates with scalar weight.

cmplt t1 t2 returns true where t1 < t2, false elsewhere.

less t1 t2 is synonym for cmplt.

less_s t scalar checks if each element is less than scalar and returns booleans.

cmpne t1 t2 returns true where t1 ≠ t2, false elsewhere.

not_equal t1 t2 is synonym for cmpne.

not_equal_s t scalar compares each element with scalar for inequality and returns booleans.

cmpeq t1 t2 returns true where t1 = t2, false elsewhere.

equal t1 t2 is synonym for cmpeq.

equal_s t scalar compares each element with scalar for equality and returns booleans.

cmpgt t1 t2 returns true where t1 > t2, false elsewhere.

greater t1 t2 is synonym for cmpgt.

greater_s t scalar checks if each element is greater than scalar and returns booleans.

cmple t1 t2 returns true where t1 ≤ t2, false elsewhere.

less_equal t1 t2 is synonym for cmple.

less_equal_s t scalar checks if each element is less than or equal to scalar and returns booleans.

cmpge t1 t2 returns true where t1 ≥ t2, false elsewhere.

greater_equal t1 t2 is synonym for cmpge.

greater_equal_s t scalar checks if each element is greater than or equal to scalar and returns booleans.

array_equal t1 t2 returns scalar 1 if all elements equal, 0 otherwise.

Broadcasts inputs before comparison. Returns 0 if shapes incompatible.

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

maximum t1 t2 returns element-wise maximum.

maximum_s t scalar returns maximum of each element and scalar.

rmaximum_s scalar t is maximum_s t scalar.

imaximum target value computes maximum in-place.

imaximum_s t scalar computes maximum with scalar in-place.

minimum t1 t2 returns element-wise minimum.

minimum_s t scalar returns minimum of each element and scalar.

rminimum_s scalar t is minimum_s t scalar.

iminimum target value computes minimum in-place.

iminimum_s t scalar computes minimum with scalar in-place.

logical_and t1 t2 computes element-wise AND.

Non-zero values are true.

logical_or t1 t2 computes element-wise OR.

logical_xor t1 t2 computes element-wise XOR.

logical_not t computes element-wise NOT.

Returns 1 - x. Non-zero values become 0, zero becomes 1.

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

isinf t returns 1 where infinite, 0 elsewhere.

Detects both positive and negative infinity. Non-float types return all 0s.

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

isnan t returns 1 where NaN, 0 elsewhere.

NaN is the only value that doesn't equal itself. Non-float types return all 0s.

  # let x = create float32 [| 3 |] [| 1.; Float.nan; Float.infinity |] in
    isnan x
  - : (bool, bool_elt) t = [false, true, false]

isfinite t returns 1 where finite, 0 elsewhere.

Finite means not inf, -inf, or NaN. Non-float types return all 1s.

  # let x = create float32 [| 4 |] [| 1.; Float.infinity; Float.nan; -0. |] in
    isfinite x
  - : (bool, bool_elt) t = [true, false, false, true]

where cond if_true if_false selects elements based on condition.

Returns if_true where cond is true, if_false elsewhere. All three inputs broadcast to common shape.

• raises Invalid_argument

  if shapes incompatible for broadcasting

  # let cond = create bool [| 3 |] [| true; false; true |] in
    let if_true = create int32 [| 3 |] [| 2l; 3l; 4l |] in
    let if_false = create int32 [| 3 |] [| 5l; 6l; 7l |] in
    where cond if_true if_false
  - : (int32, int32_elt) t = [2, 6, 4]
  # 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 limits values to range.

Elements below min become min, above max become max.

clip ?min ?max t is synonym for clamp.

bitwise_xor t1 t2 computes element-wise XOR.

bitwise_or t1 t2 computes element-wise OR.

bitwise_and t1 t2 computes element-wise AND.

bitwise_not t computes element-wise NOT.

invert t is synonym for bitwise_not.

lshift t shift left-shifts elements by shift bits.

Equivalent to multiplication by 2^shift. Overflow wraps around.

• raises Invalid_argument

  if shift negative or non-integer dtype

  # let x = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    lshift x 2
  - : (int32, int32_elt) t = [4, 8, 12]

rshift t shift right-shifts elements by shift bits.

Equivalent to integer division by 2^shift (rounds toward zero).

• raises Invalid_argument

  if shift negative or non-integer dtype

  # let x = create int32 [| 3 |] [| 8l; 9l; 10l |] in
    rshift x 2
  - : (int32, int32_elt) t = [2, 2, 2]

Infix operators

sum ?axes ?keepdims t sums elements along specified axes.

Default sums all axes (returns scalar). If keepdims is true, retains reduced dimensions with size 1. Negative axes count from end.

• raises Invalid_argument

  if any axis is out of bounds

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

max ?axes ?keepdims t finds maximum along axes.

Default reduces all axes. NaN propagates (any NaN input gives NaN output).

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

min ?axes ?keepdims t finds minimum along axes.

Default reduces all axes. NaN propagates (any NaN input gives NaN output).

  # let x = create float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in
    min x |> item []
  - : float = 1.
  # let x = create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |] in
    min ~axes:[ 0 ] x
  - : (float, float32_elt) t = [1, 2]

prod ?axes ?keepdims t computes product along axes.

Default multiplies all elements. Empty axes give 1.

  # let x = create int32 [| 3 |] [| 2l; 3l; 4l |] in
    prod x |> item []
  - : int32 = 24l
  # let x = create int32 [| 2; 2 |] [| 1l; 2l; 3l; 4l |] in
    prod ~axes:[ 0 ] x
  - : (int32, int32_elt) t = [3, 8]

cumsum ?axis t computes the inclusive cumulative sum. Defaults to flattening the tensor (row-major order) when axis is omitted.

cumprod ?axis t computes the inclusive cumulative product. Defaults to flattening the tensor when axis is omitted.

cummax ?axis t computes the inclusive cumulative maximum. NaNs propagate for floating-point dtypes. Defaults to flattening when axis is omitted.

cummin ?axis t computes the inclusive cumulative minimum. NaNs propagate for floating-point dtypes. Defaults to flattening when axis is omitted.

mean ?axes ?keepdims t computes arithmetic mean along axes.

Sum of elements divided by count. NaN propagates.

  # let x = create float32 [| 4 |] [| 1.; 2.; 3.; 4. |] in
    mean x |> item []
  - : float = 2.5
  # let x = create float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in
    mean ~axes:[ 1 ] x
  - : (float, float32_elt) t = [2, 5]

var ?axes ?keepdims ?ddof t computes variance along axes.

ddof is delta degrees of freedom. Default 0 (population variance). Use 1 for sample variance. Variance = E(X - E[X])² / (N - ddof).

• raises Invalid_argument

  if ddof >= number of elements

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

std ?axes ?keepdims ?ddof t computes standard deviation.

Square root of variance: sqrt(var(t, ddof)). See var for ddof meaning.

  # let x = create float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |] in
    std x |> item [] |> Float.round
  - : float = 1.
  # let x = create float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |] in
    std ~ddof:1 x |> item [] |> Float.round
  - : float = 2.

all ?axes ?keepdims t tests if all elements are true (non-zero).

Returns true if all elements along axes are non-zero, false otherwise.

  # let x = create int32 [| 3 |] [| 1l; 2l; 3l |] in
    all x |> item []
  - : bool = true
  # let x = create int32 [| 3 |] [| 1l; 0l; 3l |] in
    all x |> item []
  - : bool = false
  # let x = create int32 [| 2; 2 |] [| 1l; 0l; 1l; 1l |] in
    all ~axes:[ 1 ] x
  - : (bool, bool_elt) t = [false, true]

any ?axes ?keepdims t tests if any element is true (non-zero).

Returns true if any element along axes is non-zero, false if all are zero.

  # let x = create int32 [| 3 |] [| 0l; 0l; 1l |] in
    any x |> item []
  - : bool = true
  # let x = create int32 [| 3 |] [| 0l; 0l; 0l |] in
    any x |> item []
  - : bool = false
  # let x = create int32 [| 2; 2 |] [| 0l; 0l; 0l; 1l |] in
    any ~axes:[ 1 ] x
  - : (bool, bool_elt) t = [false, true]

argmax ?axis ?keepdims t finds indices of maximum values.

Returns index of first occurrence for ties. If axis not specified, operates on flattened tensor and returns scalar.

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

argmin ?axis ?keepdims t finds indices of minimum values.

Returns index of first occurrence for ties. If axis not specified, operates on flattened tensor and returns scalar.

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

sort ?descending ?axis t sorts elements along axis.

Returns (sorted_values, indices) where indices map sorted positions to original positions. Default sorts last axis in ascending order.

Algorithm: Bitonic sort (parallel-friendly, stable)

• Pads to power of 2 with inf/-inf for correctness
• O(n log² n) comparisons, O(log² n) depth
• Stable: preserves relative order of equal elements
• First occurrence wins for duplicate values

Special values:

• NaN: sorted to end (ascending) or beginning (descending)
• inf/-inf: sorted normally
• For integers: uses max/min values for padding

• raises Invalid_argument

  if axis out of bounds

  # let x = create int32 [| 5 |] [| 3l; 1l; 4l; 1l; 5l |] in
    sort x
  - : (int32, int32_elt) t * (int32, int32_elt) t =
  ([1, 1, 3, 4, 5], [1, 3, 0, 2, 4])
  # let x = create int32 [| 2; 2 |] [| 3l; 1l; 1l; 4l |] in
    sort ~descending:true ~axis:0 x
  - : (int32, int32_elt) t * (int32, int32_elt) t =
  ([[3, 4],
    [1, 1]], [[0, 1],
              [1, 0]])
  # let x = create float32 [| 4 |] [| Float.nan; 1.; 2.; Float.nan |] in
    let v, _ = sort x in
    v
  - : (float, float32_elt) t = [1, 2, nan, nan]

argsort ?descending ?axis t returns indices that would sort tensor.

Equivalent to snd (sort ?descending ?axis t). Returns indices such that taking elements at these indices yields sorted array.

For 1-D: resulti is the index of the i-th smallest element. For N-D: sorts along specified axis independently.

• raises Invalid_argument

  if axis out of bounds

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

dot a b computes generalized dot product.

Important: dot has different broadcasting behavior than matmul:

• matmul broadcasts batch dimensions
• dot does NOT broadcast; it concatenates non-contracted dimensions

For N-D × M-D arrays: dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m]) This can result in much larger output arrays than matmul.

Contracts last axis of a with:

• 1-D b: the only axis (axis 0)
• N-D b: second-to-last axis (axis -2)

Dimension rules:

• 1-D × 1-D: inner product, returns scalar
• 2-D × 2-D: matrix multiplication
• N-D × M-D: batched contraction over all but contracted axes

Supports broadcasting on batch dimensions. Result shape is concatenation of:

• Broadcasted batch dims
• Remaining dims from a (except last)
• Remaining dims from b (except contracted axis)

• raises Invalid_argument

  if contraction axes have different sizes or inputs are 0-D

  # let a = create float32 [| 2 |] [| 1.; 2. |] in
    let b = create float32 [| 2 |] [| 3.; 4. |] in
    dot a b |> item []
  - : float = 11.
  # let a = create float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |] in
    let b = create float32 [| 2; 2 |] [| 5.; 6.; 7.; 8. |] in
    dot a b
  - : (float, float32_elt) t = [[19, 22],
                                [43, 50]]
  # dot (ones float32 [| 3; 4; 5 |]) (ones float32 [| 5; 6 |]) |> shape
  - : int array = [|3; 4; 6|]
  # dot (ones float32 [| 2; 3; 4; 5 |]) (ones float32 [| 3; 5; 6 |]) |> shape
  - : int array = [|2; 3; 4; 6|]

matmul a b computes matrix multiplication with broadcasting.

Follows NumPy's @ operator semantics:

• 1-D × 1-D: inner product (returns scalar tensor)
• 1-D × N-D: treated as 1 × k @ ... × k × n → ... × n
• N-D × 1-D: treated as ... × m × k @ k × 1 → ... × m
• N-D × M-D: batched matrix multiply on last 2 dimensions

Broadcasting rules:

• All dimensions except last 2 are broadcast together
• For 1-D inputs, dimension is temporarily added then removed
• Inner dimensions must match: a.shape-1 == b.shape-2

Result shape:

• Batch dims: broadcast(a.shape:-2, b.shape:-2)
• Matrix dims: ..., a.shape[-2], b.shape[-1]
• 1-D adjustments applied after

• raises Invalid_argument

  if inputs are 0-D or inner dimensions mismatch

  # let a = create float32 [| 3 |] [| 1.; 2.; 3. |] in
    let b = create float32 [| 3 |] [| 4.; 5.; 6. |] in
    matmul a b |> item []
  - : float = 32.
  # 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]
  # let a = create float32 [| 2 |] [| 1.; 2. |] in
    let b = create float32 [| 2; 3 |] [| 3.; 4.; 5.; 6.; 7.; 8. |] in
    matmul a b
  - : (float, float32_elt) t = [15, 18, 21]
  # matmul (ones float32 [| 10; 3; 4 |]) (ones float32 [| 10; 4; 5 |]) |> shape
  - : int array = [|10; 3; 5|]
  # matmul (ones float32 [| 1; 3; 4 |]) (ones float32 [| 5; 4; 2 |]) |> shape
  - : int array = [|5; 3; 2|]

diagonal ?offset ?axis1 ?axis2 a extracts diagonal from 2-D planes.

• offset: diagonal offset (0=main, positive=above, negative=below)
• axis1, axis2: axes of 2-D planes (default: last two axes)

For 2-D array, returns 1-D array of diagonal elements. For N-D array, returns array with diagonals from each 2-D subarray.

• raises Invalid_argument

  if axis1 = axis2 or axes out of bounds