LayoutMap is a Keras API class that maps model-variable paths to tuples describing how each variable dimension is sharded or replicated across a device mesh. section 18.2.2
Assign layouts to model variables
Create a LayoutMap for a device mesh, then assign layouts by variable path: kernels use tuple(NULL, "model") to shard their second dimension, while biases use tuple("model") to distribute their dimension. section 18.2.2
layout_map <- keras$distribution$LayoutMap(device_mesh)
layout_map["sequential/dense/kernel"] <- tuple(NULL, "model")
layout_map["sequential/dense/bias"] <- tuple("model")
layout_map["sequential/dense_1/kernel"] <- tuple(NULL, "model")
layout_map["sequential/dense_1/bias"] <- tuple("model")
The R example shows LayoutMap construction and path-based assignments for Dense-layer kernels and biases. section 18.2.2
A layout specifies one entry for each variable dimension:
- NULL leaves a dimension unassigned to a mesh axis, while a named axis such as "model" controls its distribution. section 18.2.2
- Variable paths are string identifiers such as sequential/dense_1/kernel. section 18.2.2
