keras/src/distribution/distribution_lib.py/LayoutMap/_maybe_populate_device_mesh
class LayoutMap:    def _maybe_populate_device_mesh(self, layout):
        if layout.device_mesh is None and self.device_mesh is not None:
            layout.device_mesh = self.device_mesh