unified high level distribution apis across backends do not use currently under development and apis are not final currently only the jax backend has been implemented the tensorflow backend will be implemented in the future via tf dtensor api return all the available devices based on the device type note in a distributed setting global devices are returned args devicetype string one of cpu gpu or tpu defaults to gpu or tpu if available when devicetype is not provided otherwise will return the cpu devices return list of devices that are available for distribute computation initialize the distribution system for multihostprocess setting calling initialize will prepare the backend for execution on multihost gpu or tpus it should be called before any computations note that the parameters can also be injected via enviornment variables which can be better controlled by the launch script at startup time for certain backend that also rely on the enviornment variables to configure keras will properly forward them args jobaddresses string comma separated ip addresses for all the jobs that will form the whole computation cluster note that for jax backend only the address for job 0 coodinator is needed for certain runtime like cloud tpu this value can be none and the backend will figure it out with the tpu enviornment variables you can also config this value via enviornment variable kerasdistributionjobaddresses numprocesses int the number of workerprocesses that will form the whole computation cluster for certain runtime like cloud tpu this value can be none and the backend will figure it out with the tpu enviornment variables you can also configure this value via enviornment variable kerasdistributionnumprocesses processid int the id number of the current workerprocess the value should be ranged from 0 to numprocesses 1 0 will indicate the current workerprocess is the mastercoordinate job you can also configure this value via enviornment variable kerasdistributionprocessid example suppose there are two gpu processes and process 0 is running at address 10 0 0 1 1234 and process 1 is running at address 10 0 0 2 2345 to configure such cluster you can run on process 0 python keras distribute initialize jobaddresses10 0 0 1 1234 10 0 0 2 2345 numprocesses2 processid0 on process 1 python keras distribute initialize jobaddresses10 0 0 1 1234 10 0 0 2 2345 numprocesses2 processid1 or via the enviornment variables on process 0 python os environ kerasdistributionjobaddresses 10 0 0 1 1234 10 0 0 2 2345 os environkerasdistributionnumprocesses 2 os environkerasdistributionprocessid 0 keras distribute initialize on process 1 python os environ kerasdistributionjobaddresses 10 0 0 1 1234 10 0 0 2 2345 os environkerasdistributionnumprocesses 2 os environkerasdistributionprocessid 1 keras distribute initialize also note that for jax backend the jobaddresses can be further reduced to just the mastercoordinator address which is 10 0 0 1 1234 a cluster of computation devices for distributed computation this api is aligned with jax sharding mesh and tf dtensor mesh which represents the computation devices in the global context see more details in jax sharding mesh https jax readthedocs ioenlatestjax sharding htmljax sharding mesh and tf dtensor mesh https www tensorflow orgapidocspythontfexperimentaldtensormesh args shape tuple of list of integers the shape of the overall devicemesh e g 8 for a data parallel only distribution or 4 2 for a modeldata parallel distribution axisnames list of string the logical name of the each axis for the devicemesh the length of the axisnames should match to the rank of the shape the axisnames will be used to matchcreate the tensorlayout when distribute the data and variables devices optional list of devices defaults to all the available devices locally from keras distribution listdevices a layout to apply to a tensor this api is aligned with jax sharding namedsharding and tf dtensor layout see more details in jax sharding namedsharding https jax readthedocs ioenlatestjax sharding htmljax sharding namedsharding and tf dtensor layout https www tensorflow orgapidocspythontfexperimentaldtensorlayout args axes tuple of strings that should map to the axisnames in a devicemesh for any dimentions that doesn t need any sharding a none can be used a placeholder devicemesh optional devicemesh that will be used to create the layout the actual mapping of tensor to physical device is not known until the mesh is specified base class for variable distribution strategies a distribution has following key functionalities 1 distribute the model variables to a devicemesh 2 distribute the input data to a devicemesh 3 distribute an intermediate state tensor in the model it can create a context scope so that the framework to properly detect the distribution and distribute the variabledata accordingly args devicemesh a devicemesh instance retrieve the tensorlayout for the input data args datashape shape for the input data in list or tuple format returns the tensorlayout for the data which can be used by backend distributevalue to redistribute a input data retrieve the tensorlayout for the variable args variable a kerasvariable instance return the tensorlayout for the variable which can be used by backend distributevalue to redistribute a variable retrieve the tensorlayout for the intermediate tensor args path a string path for the correspoding tensor return the tensorlayout for the intermediate tensor which can be used by backend relayout to reshard the tensor could also return none context manager to make the distribution current originalscope distribution setdistributionself try yield finally setdistributionoriginalscope property def devicemeshself return self devicemesh def distributedatasetself dataset raise notimplementederror kerasexportkeras distribution dataparallel class dataparalleldistribution def initself devicemeshnone devicesnone if devicemesh self initializewithdevicemeshdevicemesh elif devices self initializemeshfromdevicesdevices else self initializemeshfromlistdevices self batchdimname self devicemesh axisnames0 those following attributes might get convert to public methods self numprocess distributionlib numprocesses self processid distributionlib processid self ismultiprocess self numprocess 1 def initializewithdevicemeshself devicemesh if not isinstancedevicemesh devicemesh raise valueerror expect mesh to be an instance of devicemesh freceived meshdevicemesh of type typedevicemesh super initdevicemesh if self devicemesh devices ndim 1 warnings warn expect the input mesh to be 1d but received mesh devices ndimd the first axis will be used for dataparallel sharding devicemesh devices ndim def initializemeshfromdevicesself devices devices np arraydevices devicemesh devicemesh shapedevices shape axisnamesdefaultbatchdimname devicesdevices super initdevicemesh def initializemeshfromlistdevicesself devices np arraylistdevices devicemesh devicemesh shapedevices shape axisnamesdefaultbatchdimname devicesdevices super initdevicemesh def getdatalayoutself datashape datashardspec none lendatashape datashardspec0 self batchdimname shard on the first dim return tensorlayoutdatashardspec self devicemesh def getvariablelayoutself variable variableshardspec none lenvariable shape return tensorlayoutvariableshardspec self devicemesh def gettensorlayoutself path for data parallel training the intermediate state is not changed return none def distributedatasetself dataset from tensorflow python data experimental ops import distribute as tfdatadistribute from keras utils moduleutils import tensorflow as tf if not isinstancedataset tf data dataset raise valueerror only tf data dataset is supported for fsharding got typedataset if not self ismultiprocess return dataset batchsize tfdatadistribute computebatchsizedataset if batchsize numpy 0 raise valueerror the batch size of the input dataset is unknown please config the batch size for the input dataset e g via dataset batchbatchsize perworkerbatchsize tfdatadistribute batchsizesforworker globalbatchsizebatchsize numworkersself numprocess numreplicasperworker1 we hard code this for now workerindexself processid distributeddataset dataset rebatchperworkerbatchsize distributeddataset tfdatadistribute autosharddataset distributeddataset numworkersself numprocess indexself processid numreplicasself numprocess return distributeddataset prefetchtf data autotune kerasexportkeras distribution modelparallel class modelparalleldistribution create a mesh with 2 devices for data parallelism and 4 devices for model parallelism create a layout map that shard the dense layer and conv2d layer variables on the last dimension based on the devicemesh this means the variables will be split across 4 devices any other variable that doesn t match any key in the layout map will be fully replicated set the global distribution or via with distribution scope with only the shape change for the device mesh the variables will be sharded across 8 devices instead of 4 which further reduces the memory footprint of variables on each of the device def initself devicemesh layoutmap batchdimnamenone super initdevicemesh self layoutmap layoutmap self batchdimname batchdimname or self devicemesh axisnames0 those following attributes might get convert to public methods self numprocess distributionlib numprocesses self processid distributionlib processid self ismultiprocess self numprocess 1 def getdatalayoutself datashape datashardspec none lendatashape datashardspec0 self batchdimname shard on the first dim return tensorlayoutdatashardspec self devicemesh def getvariablelayoutself variable variablelayout self layoutmapvariable path if variablelayout is not none return variablelayout variableshardspec none lenvariable shape return tensorlayoutvariableshardspec self devicemesh def gettensorlayoutself path return self layoutmappath def distributedatasetself dataset from tensorflow python data experimental ops import distribute as tfdatadistribute from keras utils moduleutils import tensorflow as tf if not isinstancedataset tf data dataset raise valueerror only tf data dataset is supported for fsharding got typedataset if not self ismultiprocess return dataset globalbatchsize tfdatadistribute computebatchsizedataset if globalbatchsize numpy 0 raise valueerror the batch size of the input dataset is unknown please config the batch size for the input dataset e g via dataset batchbatchsize we need to compute the the perworker batch size imagine we have 4 workers and each worker has 2 devices the device id will be wworkeriddeviceid from global perspective eg w00 w01 w10 etc for a mesh with global shape batch 4 model 2 we would like the mesh to be like below batch model w00 w01 w10 w11 v w02 w03 w12 w13 in this case the batch dim will be split acorss 2 workers in the case that global batch dim is smaller than the worker size eg batch 1 model 8 then the mesh will be like below batch model w00 w01 w10 w11 w20 w21 w30 w31 and the data batch will be fully replicated for each of the worker meshbatchdimindex self devicemesh axisnames index self batchdimname meshbatchdimsize self devicemesh shapemeshbatchdimindex todoscottzhu local device count can be a field of the mesh localdevicecount np prodself devicemesh shape self numprocess if meshbatchdimsize localdevicecount no sharding is needed in this case the worker will have the global batch size and data from the iterator will need to be replicated for model dimention return dataset prefetchtf data autotune else if meshbatchdimsize localdevicecount 0 raise valueerror the batch dimention of the mesh is not compatible with the local worker device count mesh batch fdim meshbatchdimsize and local device fcount localdevicecount numshards meshbatchdimsize localdevicecount perworkerbatchsize globalbatchsize numshards distributeddataset dataset rebatchperworkerbatchsize distributeddataset tfdatadistribute autosharddataset distributeddataset numworkersnumshards indexself processid numshards numreplicasnumshards return distributeddataset prefetchtf data autotune kerasexportkeras distribution layoutmap class layoutmapcollections abc mutablemapping def initself devicemeshnone self layoutmap collections ordereddict self devicemesh devicemesh def getitemself key if key in self layoutmap return self layoutmapkey matchingkeys for k in self layoutmap if re searchk key matchingkeys appendk if lenmatchingkeys 1 raise valueerror fpath key matches multiple layout fspecification keys matchingkeys please make sure each tensorvariable path only matches at most one layout specification key in the layoutmap elif lenmatchingkeys 1 return self layoutmapmatchingkeys0 return none def setitemself key layout if key in self layoutmap raise valueerror fkey already exist in the layoutmap with fvalue self layoutmapkey please make sure to not use duplicated keys if isinstancelayout tuple layout tensorlayoutaxeslayout devicemeshnone if not isinstancelayout tensorlayout raise valueerror flayout should be a tensorlayout type got typelayout self maybepopulatedevicemeshlayout self layoutmapkey layout def delitemself key let the dict to handle the key missing error return self layoutmap popkey def lenself return lenself layoutmap def iterself return iterself layoutmap property def devicemeshself return self devicemesh def maybepopulatedevicemeshself layout if layout devicemesh is none and self devicemesh is not none layout devicemesh self devicemesh layoutmap get doc layoutmap getitem doc kerasexportkeras distribution distributetensor def distributetensortensor layout if isinstancetensor kerastensor keras tensor is only used for building functional model and can t be used to alter layoutsharding return tensor return distributionlib distributetensortensor layout kerasexportkeras distribution distribution def distribution set the distribution as the global distribution setting args value a distribution instance return all the available devices based on the device type note in a distributed setting global devices are returned args device_type string one of cpu gpu or tpu defaults to gpu or tpu if available when device_type is not provided otherwise will return the cpu devices return list of devices that are available for distribute computation initialize the distribution system for multi host process setting calling initialize will prepare the backend for execution on multi host gpu or tpus it should be called before any computations note that the parameters can also be injected via enviornment variables which can be better controlled by the launch script at startup time for certain backend that also rely on the enviornment variables to configure keras will properly forward them args job_addresses string comma separated ip addresses for all the jobs that will form the whole computation cluster note that for jax backend only the address for job 0 coodinator is needed for certain runtime like cloud tpu this value can be none and the backend will figure it out with the tpu enviornment variables you can also config this value via enviornment variable keras_distribution_job_addresses num_processes int the number of worker processes that will form the whole computation cluster for certain runtime like cloud tpu this value can be none and the backend will figure it out with the tpu enviornment variables you can also configure this value via enviornment variable keras_distribution_num_processes process_id int the id number of the current worker process the value should be ranged from 0 to num_processes 1 0 will indicate the current worker process is the master coordinate job you can also configure this value via enviornment variable keras_distribution_process_id example suppose there are two gpu processes and process 0 is running at address 10 0 0 1 1234 and process 1 is running at address 10 0 0 2 2345 to configure such cluster you can run on process 0 python keras distribute initialize job_addresses 10 0 0 1 1234 10 0 0 2 2345 num_processes 2 process_id 0 on process 1 python keras distribute initialize job_addresses 10 0 0 1 1234 10 0 0 2 2345 num_processes 2 process_id 1 or via the enviornment variables on process 0 python os environ keras_distribution_job_addresses 10 0 0 1 1234 10 0 0 2 2345 os environ keras_distribution_num_processes 2 os environ keras_distribution_process_id 0 keras distribute initialize on process 1 python os environ keras_distribution_job_addresses 10 0 0 1 1234 10 0 0 2 2345 os environ keras_distribution_num_processes 2 os environ keras_distribution_process_id 1 keras distribute initialize also note that for jax backend the job_addresses can be further reduced to just the master coordinator address which is 10 0 0 1 1234 a cluster of computation devices for distributed computation this api is aligned with jax sharding mesh and tf dtensor mesh which represents the computation devices in the global context see more details in jax sharding mesh https jax readthedocs io en latest jax sharding html jax sharding mesh and tf dtensor mesh https www tensorflow org api_docs python tf experimental dtensor mesh args shape tuple of list of integers the shape of the overall devicemesh e g 8 for a data parallel only distribution or 4 2 for a model data parallel distribution axis_names list of string the logical name of the each axis for the devicemesh the length of the axis_names should match to the rank of the shape the axis_names will be used to match create the tensorlayout when distribute the data and variables devices optional list of devices defaults to all the available devices locally from keras distribution list_devices a layout to apply to a tensor this api is aligned with jax sharding namedsharding and tf dtensor layout see more details in jax sharding namedsharding https jax readthedocs io en latest jax sharding html jax sharding namedsharding and tf dtensor layout https www tensorflow org api_docs python tf experimental dtensor layout args axes tuple of strings that should map to the axis_names in a devicemesh for any dimentions that doesn t need any sharding a none can be used a placeholder device_mesh optional devicemesh that will be used to create the layout the actual mapping of tensor to physical device is not known until the mesh is specified base class for variable distribution strategies a distribution has following key functionalities 1 distribute the model variables to a devicemesh 2 distribute the input data to a devicemesh 3 distribute an intermediate state tensor in the model it can create a context scope so that the framework to properly detect the distribution and distribute the variable data accordingly args device_mesh a devicemesh instance retrieve the tensorlayout for the input data args data_shape shape for the input data in list or tuple format returns the tensorlayout for the data which can be used by backend distribute_value to redistribute a input data retrieve the tensorlayout for the variable args variable a kerasvariable instance return the tensorlayout for the variable which can be used by backend distribute_value to redistribute a variable retrieve the tensorlayout for the intermediate tensor args path a string path for the correspoding tensor return the tensorlayout for the intermediate tensor which can be used by backend relayout to reshard the tensor could also return none context manager to make the distribution current create a distributed dataset instance from the original user dataset args dataset the original global dataset instance only tf data dataset is supported at the moment returns a sharded tf data dataset instance which will produce data for the current local worker process distribution for data parallelism you can choose to create this instance by either specifing the device_mesh or devices arguments but not both the device_mesh argument is expected to be a devicemesh instance and is expected to be 1d only in case that the mesh has multiple axes then the first axis will be treated as the data parallel dimension and a warning will be raised when a list of devices are provided they will be used to construct a 1d mesh when both mesh and devices are absent then list_devices will be used to detect any available devices and create a 1d mesh from them args device_mesh optional devicemesh instance devices optional list of devices those following attributes might get convert to public methods shard on the first dim for data parallel training the intermediate state is not changed we hard code this for now distribution that shards model variables compare to dataparallel which replicates the variables across all devices modelparallel allows you to shard variables in addition to the input data to construct a modelparallel distribution you need to provide a devicemesh and a layoutmap 1 devicemesh contains physcial device information the axis names in the mesh will be used to map the variable and data layout 2 layoutmap contains the mapping between variable paths to their corresponding tensorlayout example python devices list_devices assume there are 8 devices create a mesh with 2 devices for data parallelism and 4 devices for model parallelism device_mesh devicemesh shape 2 4 axis_names batch model devices devices create a layout map that shard the dense layer and conv2d layer variables on the last dimension based on the device_mesh this means the variables will be split across 4 devices any other variable that doesn t match any key in the layout map will be fully replicated layout_map layoutmap device_mesh layout_map dense kernel none model layout_map dense bias model layout_map conv2d kernel none none none model layout_map conv2d bias model distribution modelparallel device_mesh device_mesh layout_map layout_map batch_dim_name batch set the global distribution or via with distribution scope set_distribution distribution model model_creation model compile model fit data you can quickly update the device mesh shape to change the sharding factor of the variables e g with only the shape change for the device mesh the variables will be sharded across 8 devices instead of 4 which further reduces the memory footprint of variables on each of the device device_mesh devicemesh shape 1 8 axis_names batch model devices devices to figure out a proper layout mapping rule for all the model variables you can first list out all the model variable paths which will be used as the key to map the variables to tensorlayout e g model create_model for v in model variables print v path args device_mesh devicemesh instance for physical device and its logical mapping layout_map layoutmap instance which map the variable path to the corresponding tensorlayout the axis names of the tensorlayout s should match to the axis names in the device_mesh or exception will be raised batch_dim_name optional string the axis name in the device_mesh that will be used to distribute data if unspecified the first axis from the device_mesh will be used those following attributes might get convert to public methods shard on the first dim we need to compute the the per worker batch size imagine we have 4 workers and each worker has 2 devices the device id will be w worker_id _ device_id from global perspective eg w0_0 w0_1 w1_0 etc for a mesh with global shape batch 4 model 2 we would like the mesh to be like below batch model w0_0 w0_1 w1_0 w1_1 v w0_2 w0_3 w1_2 w1_3 in this case the batch dim will be split acorss 2 workers in the case that global batch dim is smaller than the worker size eg batch 1 model 8 then the mesh will be like below batch model w0_0 w0_1 w1_0 w1_1 w2_0 w2_1 w3_0 w3_1 and the data batch will be fully replicated for each of the worker todo scottzhu local device count can be a field of the mesh no sharding is needed in this case the worker will have the global batch size and data from the iterator will need to be replicated for model dimention a dict like object that maps string to tensorlayout instances layoutmap uses a string as key and a tensorlayout as value there is a behavior difference between a normal python dict and this class the string key will be treated as a regex when retrieving the value see the docstring of get for more details see below for a usage example you can define the naming schema of the tensorlayout and then retrieve the corresponding tensorlayout instance in the normal case the key to query is usually the variable path which is the idenifier of the variable as shortcut tuple or list of axis names are also allowed when inserting as value and will be converted to tensorlayout python layout_map layoutmap device_mesh none layout_map dense kernel none model layout_2d layout_map dense bias model layout_1d layout_map conv2d kernel tensorlayout none none none model layout_map conv2d bias tensorlayout model layout_1d layout_1 layout_map dense_1 kernel layout_1 layout_2d layout_2 layout_map dense_1 bias layout_2 layout_1d layout_3 layout_map dense_2 kernel layout_3 layout_2d layout_4 layout_map dense_2 bias layout_4 layout_1d layout_5 layout_map my_model conv2d_123 kernel layout_5 layout_4d layout_6 layout_map my_model conv2d_123 bias layout_6 layout_1d layout_7 layout_map my_model conv3d_1 kernel layout_7 none layout_8 layout_map my_model conv3d_1 bias layout_8 none args device_mesh an optional devicemesh that can be used to populate the tensorlayout device_mesh if tensorlayout device_mesh is not set retrieves the corresponding layout by the string key when there isn t an exact match all the existing keys in the layout map will be treated as a regex and map against the input key again when there are multiple matches for the regex an valueerror will be raised returns none if there isn t any match found args key string key to query a layout returns corresponding layout based on the query insert tensorlayout to the layoutmap args key string key for the tensorlayout layout the tensorlayout as a shortcut tuple of string and none are also acceptable and will be converted to tensorlayout let the dict to handle the key missing error change the layout of a tensor value in the jit function execution args tensor a tensor to change the layout layout tensorlayout to be applied on the value returns a new value with the specified tensor layout keras tensor is only used for building functional model and can t be used to alter layout sharding retrieve the current distribution from global context set the distribution as the global distribution setting args value a distribution instance
import collections
import contextlib
import os
import re
import warnings

import numpy as np

from keras.api_export import keras_export
from keras.backend import KerasTensor
from keras.backend import distribution_lib
from keras.backend.common import global_state

DEFAULT_BATCH_DIM_NAME = "batch"
GLOBAL_ATTRIBUTE_NAME = "distribution"


@keras_export("keras.distribution.list_devices")
def list_devices(device_type=None):
    
    return distribution_lib.list_devices(device_type)


@keras_export("keras.distribution.initialize")
def initialize(job_addresses=None, num_processes=None, proceed_id=None):
    
    if (
        job_addresses is None
        and "KERAS_DISTRIBUTION_JOB_ADDRESSES" in os.environ
    ):
        job_addresses = os.environ["KERAS_DISTRIBUTION_JOB_ADDRESSES"]
    if (
        num_processes is None
        and "KERAS_DISTRIBUTION_NUM_PROCESSES" in os.environ
    ):
        num_processes = int(os.environ["KERAS_DISTRIBUTION_NUM_PROCESSES"])
    if proceed_id is None and "KERAS_DISTRIBUTION_PROCESS_ID" in os.environ:
        proceed_id = int(os.environ["KERAS_DISTRIBUTION_PROCESS_ID"])
    distribution_lib.initialize(job_addresses, num_processes, proceed_id)


@keras_export("keras.distribution.DeviceMesh")
class DeviceMesh:
    

    def __init__(
        self,
        shape,
        axis_names,
        devices=None,
    ):
        if not shape or not axis_names:
            raise ValueError(
                "Shape and axis_names cannot be empty. Received: "
                f"shape={shape}, axis_names={axis_names}"
            )

        if len(shape) != len(axis_names):
            raise ValueError(
                "Shape and axis_names should have same size. "
                f"Received: shape={shape}, axis_names={axis_names}"
            )
        if devices is None:
            devices = list_devices()
        devices = np.array(devices)
        if np.prod(shape) != np.prod(devices.shape):
            raise ValueError(
                "Shape does not match the number of devices. "
                f"Received: shape={shape}; devices.shape="
                f"{devices.shape}"
            )

        self._shape = shape
        self._axis_names = axis_names
        self._devices = np.reshape(devices, shape)

    @property
    def shape(self):
        return self._shape

    @property
    def axis_names(self):
        return self._axis_names

    @property
    def devices(self):
        return self._devices


@keras_export("keras.distribution.TensorLayout")
class TensorLayout:
    

    def __init__(self, axes, device_mesh=None):
        self._axes = tuple(axes)
        self._device_mesh = device_mesh
        self._validate_axes()

    @property
    def axes(self):
        return self._axes

    @property
    def device_mesh(self):
        return self._device_mesh

    @device_mesh.setter
    def device_mesh(self, device_mesh):
        if self._device_mesh is not None:
            raise ValueError(
                "Cannot override device mesh value. Existing "
                f"value is {self._device_mesh}"
            )
        self._device_mesh = device_mesh
        self._validate_axes()

    def _validate_axes(self):
        if self._device_mesh:
            valid_axis_names = set(self._device_mesh.axis_names)
            axis_names = set(self._axes) - set([None])
            if axis_names - valid_axis_names:
                raise ValueError(
                    "Invalid axis names for Layout. Valid axis "
                    f"names: {valid_axis_names}, Got {axis_names}"
                )


class Distribution:
    

    def __init__(self, device_mesh):
        self._device_mesh = device_mesh

    def get_data_layout(self, data_shape):
        
        raise NotImplementedError()

    def get_variable_layout(self, variable):
        
        raise NotImplementedError()

    def get_tensor_layout(self, path):
        
        raise NotImplementedError()

    @contextlib.contextmanager
    def scope(self):
        
        original_scope = distribution()
        set_distribution(self)
        try:
            yield
        finally:
            set_distribution(original_scope)

    @property
    def device_mesh(self):
        return self._device_mesh

    def distribute_dataset(self, dataset):
        
        raise NotImplementedError()


@keras_export("keras.distribution.DataParallel")
class DataParallel(Distribution):
    

    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]
        
        self._num_process = distribution_lib.num_processes()
        self._process_id = distribution_lib.process_id()
        self._is_multi_process = self._num_process > 1

    def _initialize_with_device_mesh(self, device_mesh):
        if not isinstance(device_mesh, DeviceMesh):
            raise ValueError(
                "Expect `mesh` to be an instance of `DeviceMesh`. "
                f"Received: mesh={device_mesh} (of type {type(device_mesh)})"
            )
        super().__init__(device_mesh)
        if self.device_mesh.devices.ndim != 1:
            warnings.warn(
                "Expect the input mesh to be 1D, but received "
                "mesh.devices.ndim=%d. "
                "The first axis will be used for data-parallel sharding.",
                device_mesh.devices.ndim,
            )

    def _initialize_mesh_from_devices(self, devices):
        devices = np.array(devices)
        device_mesh = DeviceMesh(
            shape=devices.shape,
            axis_names=[DEFAULT_BATCH_DIM_NAME],
            devices=devices,
        )
        super().__init__(device_mesh)

    def _initialize_mesh_from_list_devices(self):
        devices = np.array(list_devices())
        device_mesh = DeviceMesh(
            shape=devices.shape,
            axis_names=[DEFAULT_BATCH_DIM_NAME],
            devices=devices,
        )
        super().__init__(device_mesh)

    def get_data_layout(self, data_shape):
        data_shard_spec = [None] * len(data_shape)
        data_shard_spec[0] = self._batch_dim_name  
        return TensorLayout(data_shard_spec, self.device_mesh)

    def get_variable_layout(self, variable):
        variable_shard_spec = [None] * len(variable.shape)
        return TensorLayout(variable_shard_spec, self.device_mesh)

    def get_tensor_layout(self, path):
        
        return None

    def distribute_dataset(self, dataset):
        from tensorflow.python.data.experimental.ops import (
            distribute as tf_data_distribute,
        )

        from keras.utils.module_utils import tensorflow as tf

        if not isinstance(dataset, tf.data.Dataset):
            raise ValueError(
                "Only `tf.data.Dataset` is supported for "
                f"sharding, got {type(dataset)}"
            )
        if not self._is_multi_process:
            return dataset

        batch_size = tf_data_distribute.compute_batch_size(dataset)
        if batch_size.numpy() < 0:
            raise ValueError(
                "The batch size of the input dataset is "
                "unknown. Please config the batch size for "
                "the input dataset, e.g via `dataset.batch(batch_size)`"
            )
        per_worker_batch_size = tf_data_distribute.batch_sizes_for_worker(
            global_batch_size=batch_size,
            num_workers=self._num_process,
            num_replicas_per_worker=1,  
            worker_index=self._process_id,
        )
        distributed_dataset = dataset.rebatch(per_worker_batch_size)
        distributed_dataset = tf_data_distribute._AutoShardDataset(
            distributed_dataset,
            num_workers=self._num_process,
            index=self._process_id,
            num_replicas=self._num_process,
        )
        return distributed_dataset.prefetch(tf.data.AUTOTUNE)


@keras_export("keras.distribution.ModelParallel")
class ModelParallel(Distribution):
    

    def __init__(self, device_mesh, layout_map, batch_dim_name=None):
        super().__init__(device_mesh)
        self._layout_map = layout_map
        self._batch_dim_name = batch_dim_name or self.device_mesh.axis_names[0]

        
        self._num_process = distribution_lib.num_processes()
        self._process_id = distribution_lib.process_id()
        self._is_multi_process = self._num_process > 1

    def get_data_layout(self, data_shape):
        data_shard_spec = [None] * len(data_shape)
        data_shard_spec[0] = self._batch_dim_name  
        return TensorLayout(data_shard_spec, self.device_mesh)

    def get_variable_layout(self, variable):
        variable_layout = self._layout_map[variable.path]
        if variable_layout is not None:
            return variable_layout
        variable_shard_spec = [None] * len(variable.shape)
        return TensorLayout(variable_shard_spec, self.device_mesh)

    def get_tensor_layout(self, path):
        return self._layout_map[path]

    def distribute_dataset(self, dataset):
        from tensorflow.python.data.experimental.ops import (
            distribute as tf_data_distribute,
        )

        from keras.utils.module_utils import tensorflow as tf

        if not isinstance(dataset, tf.data.Dataset):
            raise ValueError(
                "Only `tf.data.Dataset` is supported for "
                f"sharding, got {type(dataset)}"
            )
        if not self._is_multi_process:
            return dataset

        global_batch_size = tf_data_distribute.compute_batch_size(dataset)
        if global_batch_size.numpy() < 0:
            raise ValueError(
                "The batch size of the input dataset is "
                "unknown. Please config the batch size for "
                "the input dataset, e.g via `dataset.batch(batch_size)`"
            )

        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        mesh_batch_dim_index = self.device_mesh.axis_names.index(
            self._batch_dim_name
        )
        mesh_batch_dim_size = self.device_mesh.shape[mesh_batch_dim_index]
        
        local_device_count = (
            np.prod(self.device_mesh.shape) // self._num_process
        )
        if mesh_batch_dim_size < local_device_count:
            
            
            
            return dataset.prefetch(tf.data.AUTOTUNE)
        else:
            if mesh_batch_dim_size % local_device_count != 0:
                raise ValueError(
                    "The Batch dimention of the mesh is not compatible "
                    "with the local worker device count. Mesh batch "
                    f"dim = {mesh_batch_dim_size} and local device "
                    f"count = {local_device_count}"
                )
            num_shards = mesh_batch_dim_size // local_device_count
            per_worker_batch_size = global_batch_size // num_shards

            distributed_dataset = dataset.rebatch(per_worker_batch_size)
            distributed_dataset = tf_data_distribute._AutoShardDataset(
                distributed_dataset,
                num_workers=num_shards,
                index=self._process_id % num_shards,
                num_replicas=num_shards,
            )
            return distributed_dataset.prefetch(tf.data.AUTOTUNE)


@keras_export("keras.distribution.LayoutMap")
class LayoutMap(collections.abc.MutableMapping):
    

    def __init__(self, device_mesh=None):
        self._layout_map = collections.OrderedDict()
        self._device_mesh = device_mesh

    def __getitem__(self, key):
        
        if key in self._layout_map:
            return self._layout_map[key]

        matching_keys = []
        for k in self._layout_map:
            if re.search(k, key):
                matching_keys.append(k)
        if len(matching_keys) > 1:
            raise ValueError(
                f"Path '{key}' matches multiple layout "
                f"specification keys: {matching_keys}. Please make "
                "sure each tensor/variable path only matches at most "
                "one layout specification key in the LayoutMap."
            )
        elif len(matching_keys) == 1:
            return self._layout_map[matching_keys[0]]
        return None

    def __setitem__(self, key, layout):
        
        if key in self._layout_map:
            raise ValueError(
                f"{key} already exist in the LayoutMap with "
                f"value {self._layout_map[key]}. Please make sure to "
                "not use duplicated keys."
            )
        if isinstance(layout, tuple):
            layout = TensorLayout(axes=layout, device_mesh=None)

        if not isinstance(layout, TensorLayout):
            raise ValueError(
                f"{layout} should be a TensorLayout type, got {type(layout)}"
            )
        self._maybe_populate_device_mesh(layout)
        self._layout_map[key] = layout

    def __delitem__(self, key):
        
        return self._layout_map.pop(key)

    def __len__(self):
        return len(self._layout_map)

    def __iter__(self):
        return iter(self._layout_map)

    @property
    def device_mesh(self):
        return self._device_mesh

    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


LayoutMap.get.__doc__ = LayoutMap.__getitem__.__doc__


@keras_export("keras.distribution.distribute_tensor")
def distribute_tensor(tensor, layout):
    
    if isinstance(tensor, KerasTensor):
        
        
        return tensor
    return distribution_lib.distribute_tensor(tensor, layout)


@keras_export("keras.distribution.distribution")
def distribution():
    
    return global_state.get_global_attribute(GLOBAL_ATTRIBUTE_NAME)


@keras_export("keras.distribution.set_distribution")
def set_distribution(value):
    
    global_state.set_global_attribute(GLOBAL_ATTRIBUTE_NAME, value)