docs menu

on this page

Module Nx_array.Layout

Layouts.

A layout maps an index (i{_0}, …, i{_k-1}), 0 <= i{_j} < d{_j}, to the element position offset + Σ i{_j}·s{_j}, counted in elements from a buffer's first byte: d{_j} are its extents, its shape, and s{_j} its strides. Positions are non-negative. An element at position p of a dtype of b bits occupies bits p·b to p·b + b - 1 of the buffer, where bit 8q + r is bit r of byte q, counted from the least significant: element p of a 4-bit dtype is the low half of byte p / 2 for an even p.

A layout is immutable and in one canonical form: an axis of extent 1 has stride 0, and a layout with no element has offset 0 and every stride 0. Two layouts of one shape that map every index to the same position are then equal. Its sizes are at most max_numel, so code that multiplies them checks nothing.

#type t

The type for layouts.

#val max_rank : int

max_rank is 32, the most axes a layout has.

#val max_numel : int

max_numel is 2{^53}. A layout's extents, the product of its extents other than 0, and the end of its span are at most max_numel. A product of any of its extents, a position times a width of at most 128 bits, and a sum of a few of these fit in an int.

#val fits : int iarray -> bool

fits s is true iff s is a layout's shape: at most max_rank extents, none negative, and those other than 0 multiplying to at most max_numel. It allocates nothing.

#val contiguous : int iarray -> t

contiguous s is the layout of shape s in C order at offset 0: element k in C order is at position k. Its shape is s itself.

Raises Invalid_argument unless fits s.

#val v : ?offset:int -> strides:int iarray -> int iarray -> t

v ~offset ~strides s is the layout of shape s with strides and offset (defaults to 0), in canonical form.

Raises Invalid_argument as contiguous does, if strides does not have s's length, if (d - 1)·|t| does not fit in an int for an axis of extent d and stride t, or, for a layout with an element, if a position is negative or the end of its span (span) exceeds max_numel.

Queries

rank, dim, stride, offset, numel, shape, strides, is_contiguous and is_distinct allocate nothing.

#val rank : t -> int

rank l is l's number of axes.

#val dim : t -> int -> int

dim l i is the extent of l's axis i.

Raises Invalid_argument unless 0 <= i < rank l.

#val stride : t -> int -> int

stride l i is the stride of l's axis i.

Raises Invalid_argument unless 0 <= i < rank l.

#val offset : t -> int

offset l is the position of l's index (0, …, 0).

#val numel : t -> int

numel l is l's number of indices, the product of its extents.

#val shape : t -> int iarray

shape l is l's extents.

#val strides : t -> int iarray

strides l is l's strides.

#val span : t -> int * int

span l is (lo, hi): every position l reaches lies in [lo, hi), with 0 <= lo, and (0, 0) if l has no element.

#val is_contiguous : t -> bool

is_contiguous l is true iff element k of l in C order is at position offset l + k.

#val is_distinct : t -> bool

is_distinct l is true iff l has no element or, with l's axes of extent above 1 ordered by |stride|, ties by axis, each stride exceeds the reach Σ (d{_j} - 1)·|s{_j}| of the axes before it. Then no two indices of l reach one position. It is true of every layout contiguous reaches by Permute, Slice and windows whose step is at least their extent dilation·(size - 1) + 1, and false of every broadcast and overlapping window.

Moving

#val move : Move.t -> t -> t option

move m l is the layout of m's result over l's positions: its index i reaches the position l reaches at m's map of i. It is None iff m is a Reshape that no strides express.

Raises Invalid_argument as Move.shape does on shape l.

#val coalesce : t array -> t array

coalesce ls is layouts of one shape that reach the positions ls reach, in the same C order of indices, with their axes of extent 1 dropped and adjacent axes merged where every layout lays them out as one run. Each has at least one axis.

Raises Invalid_argument unless ls has at least one layout, all of one shape, and Out_of_memory if host memory runs out.

Equality

#val equal : t -> t -> bool

equal l l' is true iff l and l' have the same shape, strides and offset: by the canonical form, iff they have one shape and map every index to the same position. It allocates nothing.

#val hash : t -> int

hash l is a hash of l such that equal l l' implies hash l = hash l'. It allocates nothing.

#val pp : Stdlib.Format.formatter -> t -> unit

pp formats a layout's shape, strides and offset.