Module Nx.Repr

of_array s a is the value whose bytes are a's. The value reads a's memory from then on: a write to it through any handle leaves what values over it read unspecified. The memory stays its caller's: no donation of the value writes it (donate). A place of the value may leave its copies queued, so a write to a's memory before a wait on the destination's set may reach the placed value.

Raises Invalid_argument naming Nx.Repr.of_array unless a lies on a device of s.

unsafe_of_array s a is of_array s a for an a its caller gives up: the value owns a's memory, as it owns memory nx allocated for it, so a donation of the value may write its result there (donate).

Warning. After the call, the caller must not read, write or keep a, an array over its memory or a bigarray over it, and no other value must read that memory: a donation of the value would change what they read.

Raises Invalid_argument naming Nx.Repr.unsafe_of_array unless a lies on a device of s.

array x is Some a iff x lies on one device, a its array, and None for a value on several devices. a may be strided or offset. It is for reading: writing through it changes x.

Raises Invalid_argument naming Nx.Repr.array for a dead or a traced x.

read x f is f a, a an array on the host of x's elements for the extent of f: over x's own memory where x lies on one device whose memory the host addresses (the host, a memory device, Metal), and otherwise a copy on the host, made for any placement, including a value cut across devices. It is a read: it waits for x's set and raises a failed check pending on it. a may be strided or offset.

Unlike array, it leaves x's memory unshared, so a donation of x after f returns may write its result there (donate). a's memory is claimed for the extent of f: a donation of x meanwhile computes into fresh memory. f must not write through a, and must not keep a, or a bigarray over it, after it returns. x stays reachable until f returns.

Raises Invalid_argument naming Nx.Repr.read for a dead or a traced x, or if x's memory is held exclusive, and a failed check as a wait does.

of_shards p arrays is the value whose bytes are arrays, one per device of p in order, each that device's window of one shape: along an axis p cuts, the sum of its windows' extents, and along any other, each array's.

Raises Invalid_argument naming Nx.Repr.of_shards unless there is one array per device of p, each on its device, all of one rank that has every axis p cuts, and each array is its device's window of that shape.

unsafe_of_shards p arrays is of_shards p arrays for arrays their caller gives up: the value owns their memory, as it owns memory nx allocated for it, so a donation of the value may write its result there on each device (donate).

Warning. After the call, the caller must not read, write or keep an array of arrays, an array over its memory or a bigarray over it, and no other value must read that memory: a donation of the value would change what they read.

Raises Invalid_argument naming Nx.Repr.unsafe_of_shards as of_shards raises.

shards x is x's arrays, one per device of x's placement, in order: [| a |] for a value on one device. The arrays are for reading, as array's.

Raises Invalid_argument naming Nx.Repr.shards for a dead or a traced x.