keras/src/distribution/distribution_lib.py/LayoutMap/__init__
class LayoutMap:    def __init__(self, device_mesh):
        self._layout_map = collections.OrderedDict()
        self._device_mesh = device_mesh