LayoutMap tells Keras how to shard or replicate model variables across a device mesh. It maps each variable path to a tuple with one layout entry per variable dimension, letting computation follow the variable placement. section 18.2.2
Assigning a dense layer layout
Create a LayoutMap for a mesh with a named model axis, then assign layouts such as (None, "model") to a kernel and ("model",) to its bias. Here, the kernel's second dimension and the bias are split across the model axis, while None leaves the kernel's first dimension replicated. section 18.2.2
layout_map = keras.distribution.LayoutMap(device_mesh) layout_map["sequential/dense/kernel"] = (None, "model") layout_map["sequential/dense/bias"] = ("model",) layout_map["sequential/dense_1/kernel"] = (None, "model") layout_map["sequential/dense_1/bias"] = ("model",)
Use the concrete LayoutMap construction and variable-path assignments as a starting template for model parallelism. section 18.2.2
