do not use distribution related class for jax backend this is just a prototype and we might want to unify it with other backends in the future return all the available devices based on the device type note that this should return the global devices in a distributed setting args devicetype string 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 create a distributed variable for jax since jax doesn t have a variable class this will just return a jax array with the corresponding layoutsharding specified note that this function should be used in eager context not in jitted function args value the initial value of the variable layout tensorlayout for the created variable or a jax sharding sharding instance returns jax array which is the distributed variable skip the relayout if the value is already having the proper sharding need to only distribute the value to local addressible devices and repack them back into global format distribute the tensor based on the layout note that this function can be used both in eager context or within a jitted function args tensor jax array that need to be distributed layout tensorlayout for the distribution information or a jax sharding sharding instance returns distributed value todoscottzhu this might not be a cheap check we should consider have some proper jax api for doing this check need to only distribute the value to local addressible devices and repack them back into global format distribute the input data with the corresponding layout note that the inputs here is a local worker batch within the local worker the data need to be further partitioned to map to the each of the devices args inputs jax array that is already sharded to a local process size layout tensorlayout for the distribution information or a jax sharding sharding instance returns distributed inputs thats been properly put to local devices we need the jax mesh information to determine how to place the data on to each of the worker this is data parallel mesh only we will split the full data across the batch dim datamodel parallel in this case we need to check if the mesh batch dim shape is large than number of local devices so that we can decide whether a split is needed for the data or a repeatcopy of the data is needed for each of the device todoscottzhu the mesh batch dim name is not available here since we only have jax mesh we assume the first dim is for batch and second dim is for model for now no split needed we only need to repeat here note that global batch size is not simply perprocessbatchsize numprocess it actually depends on the model dim size when user provide all the job addresses we will split and get the first one which is the coordinator do a sanity check to make sure the number of addresses also match the numprocesses return the number of processes for the current distribution setting return jax processcount def processid convert the devicemesh to jax backend specific mesh args devicemesh devicemesh instance to convert returns a jax sharding mesh instance convert the tensorlayout to jax backend specific sharding args tensorlayout tensorlayout instance to convert returns a jax sharding namedsharding instance return all the available devices based on the device type note that this should return the global devices in a distributed setting args device_type string 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 create a distributed variable for jax since jax doesn t have a variable class this will just return a jax array with the corresponding layout sharding specified note that this function should be used in eager context not in jitted function args value the initial value of the variable layout tensorlayout for the created variable or a jax sharding sharding instance returns jax array which is the distributed variable skip the relayout if the value is already having the proper sharding need to only distribute the value to local addressible devices and repack them back into global format distribute the tensor based on the layout note that this function can be used both in eager context or within a jitted function args tensor jax array that need to be distributed layout tensorlayout for the distribution information or a jax sharding sharding instance returns distributed value todo scottzhu this might not be a cheap check we should consider have some proper jax api for doing this check need to only distribute the value to local addressible devices and repack them back into global format distribute the input data with the corresponding layout note that the inputs here is a local worker batch within the local worker the data need to be further partitioned to map to the each of the devices args inputs jax array that is already sharded to a local process size layout tensorlayout for the distribution information or a jax sharding sharding instance returns distributed inputs thats been properly put to local devices we need the jax mesh information to determine how to place the data on to each of the worker this is data parallel mesh only we will split the full data across the batch dim data model parallel in this case we need to check if the mesh batch dim shape is large than number of local devices so that we can decide whether a split is needed for the data or a repeat copy of the data is needed for each of the device todo scottzhu the mesh batch dim name is not available here since we only have jax mesh we assume the first dim is for batch and second dim is for model for now no split needed we only need to repeat here note that global batch size is not simply per_process_batch_size num_process it actually depends on the model dim size when user provide all the job addresses we will split and get the first one which is the coordinator do a sanity check to make sure the number of addresses also match the num_processes return the number of processes for the current distribution setting return the current process id for the distribution setting convert the devicemesh to jax backend specific mesh args device_mesh devicemesh instance to convert returns a jax sharding mesh instance convert the tensorlayout to jax backend specific sharding args tensor_layout tensorlayout instance to convert returns a jax sharding namedsharding instance
import jax
import numpy as np

from keras.utils import jax_utils


def list_devices(device_type=None):
    
    device_type = device_type.lower() if device_type else None
    jax_devices = jax.devices(backend=device_type)
    return [f"{device.platform}:{device.id}" for device in jax_devices]


def distribute_variable(value, layout):
    
    if not isinstance(layout, jax.sharding.Sharding):
        layout = _to_jax_layout(layout)
    if isinstance(
        value, (jax.Array, jax.numpy.ndarray)
    ) and value.sharding.is_equivalent_to(layout, ndim=len(value.shape)):
        
        return value

    if layout.is_fully_addressable:
        return jax.device_put(value, layout)
    else:
        
        
        mapping = layout.addressable_devices_indices_map(value.shape)
        local_values = jax.device_put(
            [value[i] for i in mapping.values()], list(mapping.keys())
        )
        global_value = jax.make_array_from_single_device_arrays(
            value.shape, layout, local_values
        )
        return global_value


def distribute_tensor(tensor, layout):
    
    if not isinstance(layout, jax.sharding.Sharding):
        layout = _to_jax_layout(layout)
    
    
    if jax_utils.is_in_jax_tracing_scope():
        return jax.lax.with_sharding_constraint(tensor, layout)

    if layout.is_fully_addressable:
        return jax.device_put(tensor, layout)
    else:
        
        
        mapping = layout.addressable_devices_indices_map(tensor.shape)
        local_values = jax.device_put(
            [tensor[i] for i in mapping.values()], list(mapping.keys())
        )
        global_value = jax.make_array_from_single_device_arrays(
            tensor.shape, layout, local_values
        )
        return global_value


def distribute_data_input(inputs, layout):
    
    if not isinstance(layout, jax.sharding.Sharding):
        layout = _to_jax_layout(layout)
    if layout.is_fully_addressable:
        return jax.device_put(inputs, layout)

    
    
    jax_mesh = layout.mesh
    mesh_rank = len(jax_mesh.shape)
    per_process_batch_size = inputs.shape[0]
    if mesh_rank == 1:
        
        
        num_split = jax.local_device_count()
        per_replica_batch_size = per_process_batch_size // num_split
        if per_process_batch_size % per_replica_batch_size != 0:
            raise ValueError(
                f"The local batch size {per_process_batch_size} is not"
                "divisible by the number of local replicas "
                f"{num_split}"
            )
        global_batch_size = per_process_batch_size * jax.process_count()
        per_replica_batches = np.split(inputs, num_split, axis=0)
    elif mesh_rank == 2:
        
        
        
        
        
        
        
        
        mesh_batch_dim_size = list(jax_mesh.shape.values())[0]
        local_device_count = jax.local_device_count()
        if mesh_batch_dim_size < local_device_count:
            
            global_batch_size = per_process_batch_size
            per_replica_batches = [inputs for _ in range(local_device_count)]
        else:
            
            
            global_batch_size = per_process_batch_size * (
                mesh_batch_dim_size // local_device_count
            )
            per_replica_batches = np.split(inputs, local_device_count, axis=0)
    else:
        raise ValueError(
            "Only 1D or 2D mesh is supported at the moment. "
            f"Received mesh shape = {jax_mesh.shape}"
        )

    global_shape = (global_batch_size,) + inputs.shape[1:]
    global_batch_array = jax.make_array_from_single_device_arrays(
        global_shape,
        layout,
        arrays=[
            jax.device_put(batch, device)
            for batch, device in zip(
                per_replica_batches, layout.addressable_devices
            )
        ],
    )
    return global_batch_array


def initialize(job_addresses, num_processes, process_id):
    if job_addresses and "," in job_addresses:
        
        
        job_addresses = job_addresses.split(",")
        
        
        if num_processes is not None and num_processes != len(job_addresses):
            raise ValueError(
                f"The provided job_addresses {job_addresses} has "
                f"{len(job_addresses)} jobs, but num_processes is "
                f"{num_processes}"
            )
        corrdinator_address = job_addresses[0]
    else:
        corrdinator_address = job_addresses

    jax.distributed.initialize(
        corrdinator_address=corrdinator_address,
        num_processes=num_processes,
        process_id=process_id,
    )


def num_processes():
    
    return jax.process_count()


def process_id():
    
    return jax.process_index()


def _to_jax_device(device_id):
    if isinstance(device_id, jax.Device):
        return device_id
    device_type, index = device_id.split(":")
    index = int(index)
    devices = jax.devices(backend=device_type)
    if index >= len(devices):
        raise ValueError(f"Unknown device: {device_id}")
    return devices[index]


def _to_jax_mesh(device_mesh):
    
    shape = device_mesh.devices.shape
    devices = [_to_jax_device(d) for d in device_mesh.devices.flatten()]
    devices = np.array(devices).reshape(shape)
    return jax.sharding.Mesh(devices, device_mesh.axis_names)


def _to_jax_layout(tensor_layout):
    
    if tensor_layout.device_mesh is None:
        raise ValueError(
            "Cannot create sharding when device mesh is not set "
            "for TensorLayout."
        )
    partition_spec = jax.sharding.PartitionSpec(*tensor_layout.axes)
    jax_mesh = _to_jax_mesh(tensor_layout.device_mesh)
    return jax.sharding.NamedSharding(jax_mesh, partition_spec)