docs menu

on this page

Module Nx_kernel.Shape

Shapes of destinations.

Each function is the shapes of a kernel's destinations for operands of the shapes it takes, the destinations its caller allocates, or why the operands do not fit. A call of other shapes does not fit. Every shape it derives, a padded operand's and each destination's, is one a layout admits (Nx_array.Layout.max_numel), and one past it is an Error. An axis is checked against the rank it indexes, in [0, rank).

#type t = int array

The type for shapes.

#type 'a result := ('a, string) Stdlib.result

Loops

#val padded : pad -> t -> t result

padded p x is the shape a Padded load by p of an operand of shape x has. It is an Error unless p's lo, hi and interior have x's rank, interior is not negative, each padded extent is not negative, and windows are on strictly increasing axes of the padded shape, with size, step and dilation at least 1. A window longer than its axis gives it no window, as Nx_array.Move.t.Window counts them.

#val reduce : axes:int iarray -> reduction array -> t -> t result

reduce ~axes rs y is the shape of each destination of reductions rs along axes of a loop of shape y: y without axes, and y itself for no axis. It is an Error if axes are not strictly increasing axes of y, or if a Monoid Max, Monoid Min or Arg of rs has a destination element and no term.

#val scan : axis:int -> t -> t result

scan ~axis y is the shape of a scan's destinations along axis of a loop of shape y: y.

#val carry : axis:int -> t -> t result

carry ~axis y is the shape of a carry of a scan along axis of a loop of shape y: y with extent 1 along axis, with axis as scan takes it.

Gathers, scatters and sorts

#val gather : axis:int -> t -> t -> t result

gather ~axis idx x is the shape of a gather of x at positions of shape idx: idx, which has x's rank and extents but along axis.

#val gather_ranges : size:int -> t -> t -> t -> t result

gather_ranges ~size starts ends x is the shape of a range gather's destination: x with size elements along axis a = rank starts - 1. It is an Error unless starts and ends have one shape of at least one axis, x has more than a axes and starts' extents before a, and size is not negative.

#val scatter : axis:int -> t -> t -> t -> t result

scatter ~axis into idx updates is the shape of a scatter into into: into. idx and updates have one shape, into's but along axis.

#val sort : axis:int -> k:int option -> t -> t result

sort ~axis ~k x is the shape of a sort's values and positions: x with k elements along axis with Some k, x with None. It is an Error for a k past that axis's extent or negative.

Numbering

#val unique : size:int -> t -> (t * t * t) result

unique ~size keys is the shapes of a numbering's ids, of its first and counts, and of its count, for keys of [n; w]: [n], [size] and []. It is an Error unless keys has rank 2 and size is not negative.

Assemblies

#val assemble : t -> Nx_array.Move.range iarray array -> t array result

assemble y regions is the shape of each piece of an assembly of shape y: its region's, the elements regions.(i) keep of each axis of y, as Nx_array.Move.t.Slice keeps them. A fold's operand has the shape S.fold states.

Contraction

val contract : 
  batch:(int * int) iarray ->
  contracting:(int * int) iarray ->
  segments:t segments option ->
  (t * t option) ->
  (t * t option) ->
  t result

contract ~batch ~contracting ~segments (a, sa) (b, sb) is the shape of a contraction of a with b, sa and sb their scales' shapes: the batch axes in pair order, then a's free axes, then b's, in axis order. Rows drops b's b_axis; Contracting puts the groups' axis first. It is an Error unless each pair names an axis of a and an axis of b of one extent, no axis is in two pairs, each scale fits its operand (scale) and the segments' ids have the shape segments states, with a_axis and b_axis free axes, pair a contracting pair and groups at least 0.

Fourier transforms

#val c2c : axes:int iarray -> t -> t result

c2c ~axes x is x, where axes are strictly increasing axes of x, not empty.

#val r2c : axes:int iarray -> t -> t result

r2c ~axes x is x but n/2 + 1 bins along the last of axes, n x's extent there, with axes as c2c takes them.

#val c2r : axes:int iarray -> n:int -> t -> t result

c2r ~axes ~n x is x but n points along the last of axes, with axes as c2c takes them, where x has n/2 + 1 bins there and n is not negative.

Matrices

Each takes a of rank at least two, its last two axes m rows and n columns, k = min m n; its other axes, the batch, lead each shape it gives.

#val cholesky : t -> t result

cholesky a is a, square.

#val lu : t -> (t * t * t) result

lu a is the shapes of lu, pivots and perm: a, […; k] and […; m].

#val qr : factors -> t -> (t * t) result

qr f a is the shapes of q and r: […; m; k] and […; k; n] under Reduced, […; m; m] and […; m; n] under Complete.

#val svd : factors -> t -> (t * t * t) result

svd f a is the shapes of s, u and vh: […; k], then […; m; k] and […; k; n] under Reduced, […; m; m] and […; n; n] under Complete.

#val eigh : t -> (t * t) result

eigh a is the shapes of w and v of a, square: […; n] and a, for S.eigh and S.eig alike.

#val solve_triangular : t -> t -> t result

solve_triangular a b is b, […; n; r], for a square of n and b of a's batch.

Attention

#val attention : t -> (t, t) kv -> lengths:t option -> t array -> (t * t) result

attention q kv ~lengths loads is the shapes of attention's o and lse, [lead…; hq; tq; dv] and [lead…; hq; tq], for q [lead…; hq; tq; dk]. It is an Error unless q has at least three axes, kv's arrays have the shapes kv states, hk divides hq (zero heads divide zero), lengths are [lead…], each load is [lead…; hq; tq; tk] and tk is one a layout admits.