Place an array on a grid of 8 CPU devices with a named sharding, then report which block of the array each device holds.

A device mesh arranges the 8 devices into a grid of shape mesh_shape, for example (2, 4), and gives each grid axis the matching name from axis_names. Build the mesh so that the devices, taken in order of their ids, fill the grid in row-major order. Device d then sits at the grid position whose row-major flat index is d, so in a (2, 4) mesh device 5 sits in row 1 and column 1.

spec is a PartitionSpec with at most one entry per dimension of x, read from the first dimension onwards. An entry of None keeps that dimension whole on every device. An entry that names a mesh axis cuts the dimension into as many equal blocks as that axis has devices. An entry that is a tuple of names cuts it into as many blocks as all of those axes have together, and the blocks are numbered in row-major order over those axes with the first name changing slowest. Dimensions after the last entry stay whole, and devices that differ only along a mesh axis that spec never names hold identical copies. The examples write spec as a list, but the function receives a PartitionSpec such as PartitionSpec("data", None), and a nested list entry such as ["model", "data"] arrives as a tuple of names.

Return a dict with three keys. array is x placed on the mesh with a NamedSharding built from the mesh and spec, a float32 JAX array with the shape and values of x. shard_shape is an int32 JAX array of shape (8, n), where n is the number of dimensions of x, and its row d is the shape of the block that device d holds. shard_start is an int32 JAX array of shape (8, n), and its row d gives, for every dimension, the index at which device d's block begins.