A device mesh is a grid of computational devices organized along named axes. One axis can support data parallelism while another supports model parallelism. section 18.2.2
A 2 × 4 GPU mesh
An eight-GPU arrangement can use two devices along the data axis and four along the model axis. Keras then makes two copies of a model split across four GPUs, processing separate sub-batches in parallel. section 18.2.2
device_mesh <- keras$distribution$DeviceMesh(
shape = shape(2, 4), #1
axis_names = c("data", "model") #2
)
Create a named two-dimensional mesh in Keras by specifying its shape and axis names. section 18.2.2


