keras/src/distribution/distribution_lib.py/ModelParallel/__init__
class ModelParallel:    def __init__(self, *, layout_map=None, batch_dim_name=None, **kwargs):
        kwargs.pop("device_mesh", None)
        if layout_map is None:
            raise ValueError("You must specify a layout_map argument.")
        if not isinstance(layout_map, LayoutMap):
            raise ValueError(
                "Argument `layout_map` must be a `LayoutMap` instance. "
                f"Received: layout_map={layout_map}"
            )
        device_mesh = layout_map.device_mesh
        super().__init__(device_mesh)
        self._layout_map = layout_map
        self._batch_dim_name = batch_dim_name or self.device_mesh.axis_names[0]

        # Those following attributes might get convert to public methods.
        self._num_process = distribution_lib.num_processes()
        self._process_id = distribution_lib.process_id()
        self._is_multi_process = self._num_process > 1