do not use distribution related class for tensorflow 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 default 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 dtensor doesn t support getting global devices even when knowing the mesh use tf api instead to get global devices coordinator service is enabled by default with dtensor so that listlogicaldevices returns a list of global devices more context can be found in b254911601 todo convert the devicemesh to tensorflow backend specific mesh args devicemesh devicemesh instance to convert returns a tf dtensor mesh instance convert the tensorlayout to tensorflow backend specific sharding args tensorlayout tensorlayout instance to convert returns a tf dtensor layout 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 default 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 dtensor doesn t support getting global devices even when knowing the mesh use tf api instead to get global devices coordinator service is enabled by default with dtensor so that list_logical_devices returns a list of global devices more context can be found in b 254911601 todo convert the devicemesh to tensorflow backend specific mesh args device_mesh devicemesh instance to convert returns a tf dtensor mesh instance convert the tensorlayout to tensorflow backend specific sharding args tensor_layout tensorlayout instance to convert returns a tf dtensor layout instance
import tensorflow as tf
from tensorflow.experimental import dtensor


def list_devices(device_type=None):
    
    device_type = device_type.upper() if device_type else "CPU"

    
    
    
    
    tf_devices = tf.config.list_logical_devices(device_type=device_type)
    return [
        f"{device.device_type.lower()}:{device.name.split(':')[-1]}"
        for device in tf_devices
    ]


def distribute_value(value, tensor_layout):
    
    pass


def _to_dtensor_mesh(device_mesh):
    
    mesh_dims = list(zip(device_mesh.axis_names, device_mesh.shape))
    return dtensor.create_distributed_mesh(
        mesh_dims=mesh_dims, local_devices=device_mesh.devices.flatten()
    )


def _to_dtensor_layout(tensor_layout):
    
    if tensor_layout.device_mesh is None:
        raise ValueError(
            "Cannot create sharding when device mesh is not set for "
            "TensorLayout."
        )

    sharding_specs = [
        axis if axis else dtensor.UNSHARDED for axis in tensor_layout.axes
    ]
    dtensor_mesh = _to_dtensor_mesh(tensor_layout.device_mesh)
    return dtensor.Layout(sharding_specs=sharding_specs, mesh=dtensor_mesh)