Module Nx.Placement
Where a value's elements lie over its set.
type +'d tThe type for placements over a set of brand 'd.
device s d holds the whole value on d alone, a device of s; on s where s has d alone.
Raises Invalid_argument naming Nx.Placement.device unless d is a device of s (Rig.equal).
split ~axis s cuts axis into one window per device of s, in order, as evenly as its extent allows: of n elements over k devices, the first n mod k windows hold n / k + 1 and the others n / k. A window is empty only where n < k.
Raises Invalid_argument if axis is negative or not below Nx_array.Layout.max_rank.
mesh m cuts cuts each axis a of (a, names) over the mesh axes names, major first; a mesh axis no cut names holds copies. An axis cut over several mesh axes is cut as split cuts it over as many devices as their extents multiply to, numbered major first: over extents e1 and e2, the device at i1 along the first and i2 along the second holds window i1 × e2 + i2. split ~axis s is mesh (Mesh.v s [ ("x", n) ]) [ (axis, [ "x" ]) ].
Raises Invalid_argument for a name m lacks, an axis or a name in two cuts, or an axis that is negative or not below Nx_array.Layout.max_rank.
partial m sums cuts is mesh m cuts with the value a sum over the mesh axes sums. Each device holds a term, a value of the whole's shape of which it holds the window cuts gives it, and the value is the sum of the terms of its group, the devices that differ only along sums. A group's terms are numbered in row-major order over sums in the order it lists them, the first major; term 0 is its first. A product asked for its sum lists the mesh axes that cut its summed axis as the cut lists them. A value at such a placement is partial. place closes it, and a contraction asked for its sum over a summed axis cut across devices gives one (Contraction). partial m [] cuts is mesh m cuts.
Raises Invalid_argument as mesh does, and for a name in sums twice or in a cut too.
closed p is p where the value is no sum: each device holds the window p gives it of the value itself. It is p itself if p is not partial, so place (closed (placement x)) x closes x where it lies, and is x on any other placement.
equal p q is true iff p and q are of one set and every device holds the same window, and the same term of a partial value, at p and at q, whatever their brands.
same p q is Some Equal iff p and q are of one set, whatever they cut. 'd and 'e then name one brand (Device sets and placement).
uncut p is p with every axis whole on each of its devices, each mesh axis that cut one holding copies, and the value no sum. It is p itself if p cuts no axis and is not partial.