docs menu

on this page

Module Nx_kernel

Kernels: what they compute and the contract they implement.

A kernel computes one of nx's operations on arrays of one device (Rig.t). It reads its operands and writes its results into its destinations, arrays its caller allocates. Its other arguments, such as the axes it reduces, are its attributes. A kernel takes its attributes, then its destinations, then its operands, but for a scan's optional carry, which comes first, and a segmented contraction's ids, which come with its segments.

A kernel library is a module of kernels that compute on the memory of some devices. A case is an operation at some dtypes, attributes and extents, and a library declines a case when it answers without computing it. Its .mli states what it computes on, which cases it declines and the algorithms and error bounds this contract leaves to it, then includes S, so that its build checks it against the contract:

ocaml

(** Kernels for host memory. … *)

include Nx_kernel.S

A program (Prog) computes from one element of each of its operands. A loop runs a program at every index of one shape (Loops), reading each operand as a load, the array with how the loop reads it (load). A reduction, a scan, a scatter's addition, a fold (S.fold) and a contraction combine several values, the result's terms, which the operation's doc numbers. A reduction or a scan combines terms in their dtype's carrier (carrier), a dtype that holds every element of theirs. A move copies bits: a gather, a range gather, a scatter by Set, a sort, an assembly, and a map of Copy or Bitcast. A routine is a kernel of Matrices. The sections below define each operation's types and rules, then Shape, the shapes of every kernel's destinations, then S.

A call fits its kernel when its operands and destinations are as many as the kernel's doc states, of the dtypes it takes and of the shapes it and Shape state, and its attributes are ones Shape answers no Error for. nx checks that a call fits before it calls a kernel, so that an error names the arrays, then allocates the destinations. The kernel claims its destinations and operands for the call through the door, the claim of arrays Nx_array defines, and answers as Nx_array.answer states: Done, Declined for a case it does not compute, or a refusal, which names what is wrong with an operand. Any case may be Declined, within the bounds Loops states for loops. nx raises a refusal as Nx_array.refused does, and after a decline raises Invalid_argument and computes nothing in the library's place.

The caller must pass a call that fits. Where a kernel answers Declined or a refusal for a call that fits, it does so before any write or queued work. For one that does not, it answers a refusal or Declined, never Done, and reads and writes nothing outside its arrays; the library chooses which refusal, and the destinations' contents are unspecified. nx names a misfit with the refusal map_refusal, reduce_refusal, scan_refusal or attention_refusal gives, or Shape_mismatch where Shape answers Error.

A failure that depends on a float's or a complex number's value is a NaN result. A kernel raises Out_of_memory if host memory runs out, Rig.Out_of_memory if its device's does, and Rig.Lost as the door does. After a raise, the destinations' contents are unspecified.

Unless a function's doc says otherwise:

The caller must pass destinations that are C-contiguous, reach no element twice and are on the operands' device. Operands may be strided, broadcast or offset.
A destination shares no byte with an operand, except one it is identical to that the kernel reads only at the result's own index: a Plain load of map, contract's init, scatter's into, or assemble's first piece where its region is the whole result. A kernel reads such an operand at an index before writing there, so the result equals what a fresh destination would receive. The door refuses other sharing as Overlapping. It accepts a destination identical to an operand read at other indices, and the caller must not pass one; the result is then unspecified unless the library's doc states it.
A kernel is deterministic. For the same operand values, shapes, layouts and dtypes and the same attributes, it gives the same bits on every run, on any number of threads, on devices of one model, for a host library one architecture.
A result whose bits this contract fixes, an integer, a boolean, an extreme, a position, a move's or a map of kinds Prog defines exactly, depends only on the operand elements its own index reads. So a box of the results' indices, a range of them along each axis, computed alone from the parts of the operands it reads, has the bits of the whole.
Only a move specifies a NaN result's bits. Every other kernel's NaN result follows the NaN rule (Stores): its sign and payload are unspecified, and it is signalling only where an operand is a signalling NaN.

#module Prog : sig ... end

Scalar programs: what a loop computes at each index.

#type ('v, 's) a := ('v, 's) Nx_array.t

#type index := (int64, Nx_array.Dtype.int64_elt) Nx_array.t

#type 'c dtypes = {

#arrays : 'c -> int;
#dtype : 'c -> int -> Nx_array.Dtype.any;

}

The type for reading the dtypes of a collection of arrays of type 'c in place: arrays c arrays, dtype c i the dtype of the ith. A rule that decides a kernel's refusals, as attention_refusal, reads its caller's arrays through one, which the caller defines once per collection type, so that checking a call copies nothing.

Loops

#type pad = {

#lo : int iarray;
#hi : int iarray;
#interior : int iarray;
#windows : Nx_array.Move.window iarray;

}

The type for paddings of an operand of rank r, each array of lo, hi and interior of length r. Along axis i it puts lo.(i) elements before the operand's, hi.(i) after and interior.(i) between neighbours, and a negative lo or hi crops instead. Then it takes the windows of the padded operand that Nx_array.Move.t.Window windows takes. Shape.padded states which paddings an operand takes.

#type load =

#| Plain : ('v, 's) a -> load
(*
The operand through its layout.
*)
#| Padded : {
1. #x : ('v, 's) a;
2. #fill : 'v;
3. #pad : pad;
} -> load
(*
x padded with fill as pad says. The loop's shape is the padded one.
*)

The type for loads. A load is an operand of a loop together with how the loop reads it.

val map_refusal : 
  Prog.t ->
  rank:int ->
  dsts:'d dtypes ->
  'd ->
  loads:'l dtypes ->
  'l ->
  (Nx_array.answer * string) option

map_refusal p ~rank ~dsts:dv dsts ~loads:lv loads is Some (r, why) where S.map's program p over loads into dsts, read by lv and dv, on a loop of rank rank, does not fit: r the refusal that names the misfit and why what does not fit, in words; None where all fit. It reads no shape but the rank: Shape.padded gives a padded load's.

Reductions and scans

A reduction's terms are the values of one program output along the reduced axes, numbered from 0 in C order of those axes' indices.

A float result is a NaN where a term it reduces is a NaN, or where its arithmetic makes one from numbers, as +∞ plus -∞ does. To every reduction but Sum, a complex number with a NaN part is a NaN. A complex Sum adds each part as a float sum.

A float Sum or Prod reduction adds or multiplies each output's terms in the association its library states, the order and grouping of its operations, a narrow float's terms in float32, and a complex Prod multiplies them by Prog.binary.Mul. Each result rounds once into its destination, a float as Nx_array.Dtype.of_float rounds it.

#type monoid =

#| Sum
#| Prod
#| Max
#| Min
(*
The type for monoids a reduction folds with.
*)

#type extreme =

#| Max
#| Min
(*
The type for extremes a reduction finds with their positions.
*)

#val carrier : Nx_array.Dtype.any -> Nx_array.Dtype.any

carrier dt is the dtype kernels reduce and scan terms of dt in, and the dtype of a scan's carry, the state a scan continues from (S.scan): float32 for a float narrower than it, int8 and uint8 for int4 and uint4, bool for bit, else dt.

#type any_float =

#| Any_float : (float, 's) Nx_array.Dtype.t -> any_float

The type for a float dtype, whatever its storage.

#val float_carrier : (float, 's) Nx_array.Dtype.t -> any_float

float_carrier dt is carrier of the float dt, the compute type in which a float kind of dt computes before it rounds once to dt (Prog).

#type reduction =

#| Monoid of monoid
(*
Monoid Sum adds the terms from +0, as Prog.binary.Add does, so a sum of -0 terms is +0. Monoid Prod multiplies them from the first, never by 1, which would change a complex term's signed zeros or make NaN parts from its infinite ones, and an empty product is 1. Integers wrap. Monoid Max and Monoid Min are Prog.binary.Maximum and Prog.binary.Minimum of the terms, which put -0 below +0 and are any and all on booleans.
*)
#| Logsumexp
(*
Logsumexp is m + log Σ exp (x - m), m the terms' maximum. It is -∞ for no term or terms all -∞, +∞ where a term is +∞ and none is NaN, and a NaN where a term is. Each kernel library states its algorithm and its bound.
*)
#| Moments
(*
Moments is the terms' population mean, then their variance; NaN for no term or where a term is a NaN. With an infinite term and no NaN term, the mean is the sum of the infinities, a NaN where both signs occur, and the variance a NaN. The variance is never negative. Each kernel library states its algorithm and their bounds.
*)
#| Arg of extreme
(*
Arg e is the extreme e of the terms, as Monoid Max or Monoid Min finds it, then its position as an int64, that of the first term with its bits, or of the first NaN term where it is a NaN.
*)

The type for what a reduction computes from its terms.

#val accepts : reduction -> ('v, 's) Nx_array.Dtype.t -> bool

accepts r dt is true iff r reduces terms of dt: Sum and Prod every dtype but booleans, Max, Min and Arg every dtype, Logsumexp and Moments floats.

#type 'c reductions = {

#reductions : 'c -> int;
#kind : 'c -> int -> reduction;
#output : 'c -> int -> int;

}

The type for reading a collection 'c of reductions in place, as dtypes reads arrays: reductions c reductions, the ith kind c i of the program's output output c i.

val reduce_refusal : 
  Prog.t ->
  rank:int ->
  reductions:'r reductions ->
  'r ->
  dsts:'d dtypes ->
  'd ->
  loads:'l dtypes ->
  'l ->
  (Nx_array.answer * string) option

reduce_refusal p ~rank ~reductions:rv rs ~dsts:dv dsts ~loads:lv loads is as map_refusal for S.reduce's reductions rs of p, read by rv. It reads no shape but the rank: Shape.reduce gives the destinations'.

val scan_refusal : 
  Prog.t ->
  rank:int ->
  reduction ->
  output:int ->
  dsts:'d dtypes ->
  'd ->
  ?carry:Nx_array.Dtype.any ->
  loads:'l dtypes ->
  'l ->
  (Nx_array.answer * string) option

scan_refusal p ~rank r ~output:k ~dsts:dv dsts ?carry ~loads:lv loads is as map_refusal for S.scan's scan r of output k of p from a carry of dtype carry. It reads no shape but the rank: Shape.carry gives the carry's.

Error bounds

Within the bounds this contract and its library state, a float or complex result may differ between libraries, device models and layouts, and the same terms in a call of other extents may round differently. Let u be the unit roundoff of the dtype each addition or multiplication rounds to, its parts' for a complex one, and γ(k) = k u / (1 - k u). Barring overflow and underflow, and before the rounding into the destination, which adds its own, a float sum of n terms, a reduction's, a scan's or a fold's, is within γ(n - 1) Σ|x| of the exact sum, and a float product within a relative γ(n - 1). A complex product is within a normwise relative γ(2n - 2), where each Prog.binary.Mul errs by at most 2u.

Barring overflow and underflow, an output of a float or complex contraction (S.contract) with K products a·b is within e(K) (|init| + Σ|a||b|) of the exact sum of its products and init as the accumulator holds them, a complex one in modulus, |init| being 0 without init. e is the library's error factor, a function of K that the library states and that may vary with the dtypes and the device model. It is at most γ(2K + 2), u the accumulator's unit roundoff, its parts' for a complex one.

Scans

A float or complex Sum or Prod scan combines the terms up to each result by Prog.binary.Add or Prog.binary.Mul in the carrier (carrier), in the association its library states, a product from its first term, so that its first result is that term. A Logsumexp or Moments scan's result has the special values of that reduction of the terms up to it (reduction), and each kernel library states its algorithm and its bounds. Each result rounds once into its destination. A scan's bounds are on the result in the carrier, before the rounding into the destination, which adds its own. They are to first order in the carrier's unit roundoff u, barring overflow.

The result at an index is a function of the terms up to it alone, in the sense that the scan of a prefix of the axis is the prefix of the scan, bit for bit. A scan's last result need not equal the reduction of the same terms. A scan of integers or booleans, and a Max or Min scan, is exact, so scanning a slice from any term, with the carrier value of the result before it as the carry, gives the whole scan's results from that term on, bit for bit but for a NaN result, which may be any NaN.

The odd/even network

The odd/even network scans a slice by combining states, one per term. Neighbours 2j and 2j + 1 combine into pairs, the pairs scan by the same network, result 0 is state 0, result 2j + 1 is the pairs' result j, and result 2j, for j > 0, is the pairs' result j - 1 combined with state 2j. The earlier state comes first in each combination, and result i passes through at most D = max 0 (2⌈log₂ (i + 1)⌉ - 1) of them. A library that scans by it states its bounds in D.

Scatters

#type combine =

#| Set
#| Add
#| Max
#| Min
(*
The type for how a scatter combines an element with its target's.
*)

Contraction

A contraction of a with b pairs axes of a with axes of b of one extent. The result keeps the batch pairs and sums over the contracting pairs. An operand's other axes are its free axes, a's the result's rows and b's its columns. An output's products are those of a's and b's elements at its batch, row and column indices, one per index of the contracting pairs. The contracted extent is the contracting pairs' extents multiplied. The products are the output's terms, numbered in C order of the contracting pairs' indices. A product is fused into the addition that takes it. A complex contraction sums each part of an output as a float sum of products of parts. With a = ar + i·ai and b = br + i·bi, its product p's terms are numbered 2p and 2p + 1, ar·br and ai·(-bi) in the real part, ar·bi and ai·br in the imaginary part.

An output without products, where the contracted extent is 0 or its segments give it none, is its init element as the accumulator holds it, rounded into the destination, or +0 without init. No addition touches it, so an init of -0 stays -0 and a NaN init keeps its bits.

#val accumulates : ('v, 's) Nx_array.Dtype.t -> bool

accumulates dt is true iff a contraction computes in dt: every dtype but booleans and floats narrower than float32.

#type scale =

#| Exponent of (int, Nx_array.Dtype.int8_unsigned_elt) Nx_array.t
(*
Code e scales its scale block by 2^(e - 127); 255 makes it NaN.
*)
#| Value : ('v, 's) Nx_array.t -> scale
(*
Each scale read in its own dtype into the accumulator.
*)

The type for a block scale of a contraction operand's array: an array of its rank, each extent dividing the operand's, 0 only along an empty axis. The operand's element at index i is in the scale block i divided by the ratio of the extents, axis by axis. The kernel applies a scale by the operand's own index, whatever its layout. The caller must not pass the array's scale for a view that cuts an axis within its scale blocks.

#type operand = {

#array : Nx_array.any;
#scale : scale option;

}

The type for a contraction's operand: an array, through any layout (strided, broadcast, windowed), never padded, and its scale. Its element at an index is the array's times its scale block's scale.

#type 'i segments =

#| Rows of {
1. #ids : 'i;
2. #a_axis : int;
3. #b_axis : int;
}
(*
a's free axis a_axis is segmented; b's free axis b_axis, of extent g, leaves the result. Index r along a_axis multiplies b's index ids.{r} along b_axis; an output whose index along a_axis is in no group has no products.
*)
#| Contracting of {
1. #ids : 'i;
2. #groups : int;
3. #pair : int;
}
(*
Contracting pair pair is segmented on both sides into groups groups, g. The result gains a leading axis of groups, slice e summing group e's indices alone, its products numbered by their order among the group's indices from 0, its extent the group's count; a group of no index has no products.
*)

The type for a contraction cut into groups by ids, an int64 array for S.contract (its shape for Shape.contract) of the batch extents then the segmented axis's. At each batch index, the segmented axis's index r is in group ids.{r} where that lies in [0, g), and in no group otherwise.

Fourier transforms

#type direction =

#| Forward
#| Inverse

The type for a transform's direction. Along an axis of n points, Forward is y[k] = Σ_j x[j] e^(-2πi jk/n) and Inverse the same sum with e^(+2πi jk/n). Neither divides by n.

val same_precision : 
  ('v, 's) Nx_array.Dtype.t ->
  ('w, 't) Nx_array.Dtype.t ->
  bool

same_precision x y is true iff x and y have one precision, each float32 or complex64, or each float64 or complex128. A real transform, S.svd, S.eigh and S.eig pair dtypes of one precision.

Matrices

A routine computes on each matrix of a, its last two axes, of m rows and n columns, k = min m n. Each kernel library states its error bounds and the choices the routines leave open: eig's order of eigenvalues; qr's factors, unique only up to a unit factor per column of q and row of r where r's diagonal has no zero, and otherwise from the first zero on up to an orthonormal completion; the vectors of svd, eigh and eig, unique only up to a unit factor where their value is simple and up to a unitary basis of its space where it repeats; and the columns complete factors add, any orthonormal completion.

#type triangle =

#| Lower
#| Upper

The type for a square matrix's triangles, each with its diagonal.

#type factors =

#| Reduced
#| Complete

The type for how much of a matrix's unitary factors Shape.qr and Shape.svd shape: k columns of the left one and k rows of the right one, or the square ones, of m and n. A kernel reads which from its destinations' shapes. Where a destination's two shapes coincide, as q's, r's and u's do for m <= n and vh's for m >= n, they name one result. qr's factors are then the same either way, and svd reads the choice from its other factor.

#val full_precision : ('v, 's) Nx_array.Dtype.t -> bool

full_precision dt is true iff dt is float32, float64, complex64 or complex128, the dtypes the matrix routines, S.cholesky to S.solve_triangular, take.

Attention

Attention reads query rows q [lead…; hq; tq; dk] against keys and values of hk key heads, hk dividing hq, query head h reading key head h / (hq / hk). Its loop is [lead…; hq; tq; tk], a pair at each index: a query row and a key. It computes in an accumulator, float32 or float64. A pair's score is s × q·k, s the scale rounded to the accumulator and q·k the sum of the row's and the key's dk products, each element widened to the accumulator, in an order its library states. Barring overflow and underflow, q·k is within e(dk) Σ|qᵢkᵢ| of the exact sum, e the library's error factor (Error bounds), at most γ(2dk + 2). A program over the pair computes two outputs from its operand 0, the score, and its loads, operands 1 on, each of the loop's shape: output 0, a bool, which reads no score, and output 1, in the accumulator. Its Coord nodes read the loop's axes: rank - 2 the query row, rank - 1 the key.

A pair is visible where its key is before its lead index's length, tk without lengths, where a paged key's table entry is in [0, pages), and where output 0 holds. A row's lse is the log-sum-exp of output 1 over its visible pairs, and its o the sum of their values weighted by exp (x - lse), x each one's output 1, a sum that starts from +0, so a sum of zeros is +0. A row of no visible pair, or of only -∞ outputs, has o = +0 and lse = -∞. A visible pair's value counts where its weight underflows to 0: a NaN there makes the row's o a NaN. A value whose key is past the row's length, or whose table entry is outside [0, pages), is never read. Whether a NaN or an infinity at another hidden pair reaches o is the library's to state.

Barring overflow and underflow, at a row with a visible pair, let L be the exact log-sum-exp of its visible outputs 1, b a bound on their magnitudes, u the accumulator's unit roundoff and u' = u / (1 - u). Each library states a bound ε on |lse - L|. Weights are per element of o: a library may compute and round each element's weight for a pair apart. Each weight w it computes in the accumulator is exp (x - L + δ), with |δ| ≤ (1 + u) ε + (2b + ln tk) u + 5u'. Where the values are float16 or bfloat16, it may apply r(c w) / c in place of w, the quotient exact, r the rounding to nearest, ties to even, into their dtype, and c a factor it chooses with c ≥ 1 and c w ≤ 1. Before its rounding into its dtype, an element of o sums the values vⱼ of the row's visible pairs times the weights wⱼ it applies to them at that element, in an order its library states, within e(tk) Σ|wⱼ||vⱼ| of Σ wⱼ vⱼ, e as for q·k.

#type ('a, 'i) kv =

#| Dense of {
1. #k : 'a;
2. #v : 'a;
}
#| Paged of {
1. #k : 'a;
2. #v : 'a;
3. #table : 'i;
}

The type for keys and values. Dense holds k [lead…; hk; tk; dk] and v [lead…; hk; tk; dv]. Paged holds pools k [pages; hk; page; dk] and v [pages; hk; page; dv] read through table [lead…; max_pages]: key t of lead index b is the pools' entry table[b; t / page] at t mod page, and tk = max_pages · page.

val attention_refusal : 
  Prog.t ->
  q:('v, 's) Nx_array.Dtype.t ->
  rank:int ->
  o:('w, 'r) Nx_array.Dtype.t ->
  acc:(float, 'c) Nx_array.Dtype.t ->
  loads:'l dtypes ->
  'l ->
  (Nx_array.answer * string) option

attention_refusal p ~q ~rank ~o ~acc ~loads:v loads is as map_refusal for S.attention's program p over loads, read by v, a query of dtype q and rank rank, o of dtype o and lse of dtype acc. It reads no shape but the rank: Shape.attention gives the shapes.

Shapes

#module Shape : sig ... end

Shapes of destinations.

The contract

#module type S = sig ... end

The type for kernel libraries. A library's kernels keep the rules stated above and those of their sections. Whether one declines depends only on the dtypes, the attributes, the device model and the extents of the axes along which one result element reads every index: a reduction's, a scan's, a sort's or a transform's axes, a numbering's keys, the contracted extent, a routine's matrices, a gather's or a scatter's axis, a range gather's runs, attention's keys and their two dimensions, and the axes a padded load or a fold pads or windows. Each library states whether several domains may call its kernels concurrently.