Module type Nx_kernel.S

name names the kernels in messages, as "nx.cpu".

computes_on d is true iff these kernels compute on d's memory.

map p ~dsts loads stores p's output k at each index into dsts.(k), of that output's dtype. It is the one elementwise kernel. One kind over operands comes as the program of one node that Nx_kernel.Prog.of_node makes, which Nx_kernel.Prog.single and Nx_kernel.Prog.is_single recognise.

reduce p ~shape ~axes rs ~dsts loads reduces, for each (r, k) of rs, output k of p over the loop of shape shape along axes by r. Its destinations follow rs: one for a Monoid or a Logsumexp, the mean then the variance for Moments, the extreme then its int64 position for Arg. Each has the shape Nx_kernel.Shape.reduce gives of shape.

scan ?carry p ~axis (r, k) ~dsts loads stores at each index the reduction r of output k of p over the indices along axis up to that index's, inclusive (Scans). Its destinations are as reduce's, each of the loop's shape (Nx_kernel.Shape.scan).

With carry, an array of the shape Nx_kernel.Shape.carry gives, of the output's carrier, a Monoid scan of each slice continues from the slice's element of carry. Each result is that element combined with the terms up to it, and a float or complex one is within the bound of a reduction of the element and those terms (Error bounds). A NaN carry makes every result a NaN. Only a Monoid scan takes a carry.

gather ~axis ~dst idx x stores into dst at each index j x's element at j with axis axis replaced by idx's element at j. A position outside [0, d), d x's extent along axis, reads the element of zero bits: +0, false, 0 + 0i.

gather_ranges ~dst ~starts ~ends x packs runs of x along its axis a = rank starts - 1 into dst. starts and ends are [lead…; k], x is [lead…; n; cells…] and dst [lead…; size; cells…] (Nx_kernel.Shape.gather_ranges). At each lead index, position p of dst along a belongs to the first run j whose end ends.{j} exceeds p, and reads x at starts.{j} + p - e, e being ends.{j - 1}, or 0 for j = 0, computed in int64 modulo 2{^64}: it reads what the exact sum would, but where starts.{j} and e lie near opposite extremes of int64. A position past every end, or outside [0, n), reads the element of zero bits. So for ends the running sums of lengths that are not negative, run j copies its ends.{j} - e elements from starts.{j}, cut at size. It is a move: dst is gather of x at these positions.

scatter c ~unique ~axis ~dst ~into idx updates stores into dst into with updates combined by c at the positions idx holds. Update j targets into's element at j with axis axis replaced by idx's element at j; a position outside [0, d), d into's extent along axis, drops it. Where no update lands, the result is into's element, its bits unchanged, a signalling NaN's included.

Under unique, the caller promises each target at most one update, which a library may rely on: where two share one, the result there is unspecified. Add takes the dtypes Nx_kernel.Prog.binary.Add takes.

sort ~axis ~descending ~values ~positions x orders each slice of x along axis, stably, into values, bits unchanged, and their int64 positions along axis into positions. Where values is shorter than x along axis, it keeps the first elements of the sorted order. The order is the one Nx_kernel.Prog's domains state, -0 below +0, with every NaN, and every complex number with a NaN part, above +∞ and equal to each other: -∞ < … < -0 < +0 < … < +∞ < NaN. descending reverses the order and keeps equal elements in increasing position.

unique ~ids ~first ~counts ~count keys numbers the rows of keys, an [n; w] array whose elements are read as unsigned 64-bit words: two rows are one key iff their words are equal, and a w of 0 makes every row one key. Keys are numbered from 0 in order of their first rows, k of them. ids, of n entries, gets each row's key. first and counts have one length m: entry g below min m k gets key g's first row and its number of rows, and entries from k on get n and 0. count, 0-d, gets k.

assemble ~fill ~dst pieces stores into dst, at each index, the element at that index of the last piece whose region holds it, or fill where none does. Each piece has the shape of its region (Nx_kernel.Shape.assemble). Bits are copied, NaN payloads included.

fold p ~dst x stores into dst the adjoint of a Padded load by p of dst's shape, where x has the shape that load gives. dst's element at an index is the sum by Nx_kernel.Prog.binary.Add of its terms, the elements of x the load reads from that index, with no +0 before the first, so that a single -0 keeps its sign, and +0 where it reads none; in the association its library states and within the bound Error bounds states, each part of a complex sum such a sum. The library states whether a narrow float's sum rounds at each addition or once from float32. It takes the dtypes Nx_kernel.Prog.binary.Add takes.

contract ~batch ~contracting ~segments ~acc ~dst ~init a b stores into dst the contraction of a with b, of the shape Nx_kernel.Shape.contract gives:

dst[B, I, J] = init[B, I, J] + Σ_K a[B, I, K] · b[B, K, J]

computed in acc, summed in the library's order and rounded once into dst, each operand's elements times their scales, each group of segments as Nx_kernel.segments states. init has dst's shape and is one term of the sum, in the library's association; without init the sum starts from +0. An output without products is as Contraction states. Integers wrap in acc. a, b, init and dst may be of any dtypes: an element enters acc, and the sum enters dst, as Nx_kernel.Prog.op1.Cast stores it. It takes an acc that Nx_kernel.accumulates accepts.

c2c d ~axes ~dst x stores the transform d of x into dst.

r2c ~axes ~dst x stores the Forward transform of x, bins 0 to n/2 along the last of axes, n x's extent there. For n = 0 it is one bin of +0. x and dst are of one precision (Nx_kernel.same_precision).

c2r ~axes ~dst x reads n/2 + 1 bins X along the last of axes, n dst's extent there. It is the Inverse transform along each of axes but the last, then along the last the real parts of the Inverse transform of X', X'[k] X[k] for k <= n/2 and conj X[n - k] above. The imaginary parts of bin 0 and, for an even n, of bin n/2 never contribute. It is this linear map for every operand, Hermitian or not. dst and x are of one precision (Nx_kernel.same_precision).

cholesky t ~dst a stores L with a = L Lᴴ under Lower, or U = Lᴴ under Upper, +0 in the other triangle and its diagonal real and positive. It reads a's lower triangle and the real parts of its diagonal alone, whatever t. A matrix whose factorisation meets a pivot that is not positive, or NaN, is NaN.

lu ~lu ~pivots ~perm a stores into lu the unit lower triangular L below its diagonal, its ones unstored, and the upper triangular U on and above it. At step j, rows j and pivots[j] were interchanged, pivots[j] the row a search of column j down from row j ends at. The search starts at j and moves to a later row only where that row's magnitude, |re| + |im| on complex, is greater than the current one's. A NaN is never greater, and nothing is greater than a NaN. Row i of L U is row perm[i] of a. A zero pivot stays in U and its column of L is not divided.

qr ~q ~r a stores q with orthonormal columns and r, upper triangular with +0 below its diagonal, a = q r.

svd ~s ~vectors:(~u, ~vh) a stores a = u diag s vh, s descending and not negative, a zero one +0. Without vectors it stores s alone. A matrix that holds NaN or an infinity, or on which the iteration does not converge, is NaN in every result. s is of a's precision (Nx_kernel.same_precision).

eigh ~w ~v a stores the eigenvalues of a, square and Hermitian as its lower triangle and the real parts of its diagonal name it, which alone are read, ascending into w, and orthonormal columns v with a v = v diag w. Without v it stores w alone. NaN as for svd, over what is read. w is of a's precision (Nx_kernel.same_precision).

eig ~w ~v a stores the eigenvalues of a, square, into w, and columns v of unit norm with a v = v diag w. Without v it stores w alone. NaN as for svd. w and v are of a's precision (Nx_kernel.same_precision).

solve_triangular t ~adjoint ~unit_diagonal ~dst a b stores x with a x = b, or aᴴ x = b under adjoint. It reads a's triangle t, without its diagonal under unit_diagonal, whose ones it takes instead. A zero on a diagonal it reads makes that matrix's x NaN.

attention p ~scale ~o ~lse ?lengths q kv loads stores the attention of q's rows over kv (Attention), p the pair's program over its loads loads, into o and lse, whose dtype is the accumulator. lengths, [lead…], bounds each lead index's keys. The scores, their maxima and sums, the weights and lse compute in the accumulator, save a weight applied through a rounding to the values' dtype (Attention), and each element of o rounds once into o's dtype.

q, k and v are of one dtype, float16, bfloat16, float32 or float64, no wider than the accumulator, o is of the accumulator or of q's dtype, and each load is of its operand's dtype. p has one operand more than loads and two outputs, its operand 0 and output 1 are in the accumulator, its output 0 is a bool that does not read operand 0, and its Coord nodes read axes below q's rank. Nx_kernel.attention_refusal decides whether these fit. Each library states its orders and its bounds.