keras/src/distribution/distribution_lib.py/DataParallel/__init__
class DataParallel:    def __init__(self, device_mesh=None, devices=None):
        if device_mesh:
            self._initialize_with_device_mesh(device_mesh)
        elif devices:
            self._initialize_mesh_from_devices(devices)
        else:
            self._initialize_mesh_from_list_devices()

        self._batch_dim_name = 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