Модуль: tf.experimental.dtensor

Общедоступный API для пространства имён tf.experimental.dtensor.

Классы

class DTensorCheckpoint: Управляет сохранением/восстановлением отслеживаемых значений на диск для DTensor.

class DVariable: Замена для tf.Variable, которая следует размещению начального значения.

class Layout: Представляет информацию о расположении DTensor.

class Mesh: Представляет конфигурацию сетки над определённым списком измерений сетки.

Функции

call_with_layout(...): Вызывает функцию в области устройства DTensor, если layout не равно None.

check_layout(...): Утверждает, что расположение DTensor равно layout.

client_id(...): Возвращает идентификатор текущего клиента.

copy_to_mesh(...): Копирует tf.Tensor на устройство DTensor с заданным расположением.

create_distributed_mesh(...): Создаёт одно- или многоклиентскую сетку.

create_mesh(...): Создаёт одноклиентскую сетку.

device_name(...): Возвращает имя единственного устройства DTensor.

enable_save_as_bf16(...): Позволяет сохранять и восстанавливать float32 DVariables как bfloat16.

fetch_layout(...): Получает расположение DTensor.

full_job_name(...): Возвращает полное имя задачи TF для текущей или другой задачи.

heartbeat_enabled(...): Возвращает true, если служба DTensor heartbeat включена.

initialize_multi_client(...): Инициализирует DTensor для нескольких клиентов.

initialize_tpu_system(...): Инициализирует устройства TPU.

job_name(...): Возвращает имя задачи, используемое всеми клиентами в кластере DTensor.

jobs(...): Возвращает список имён задач всех клиентов в кластере DTensor.

local_devices(...): Возвращает список спецификаций устройств типа device_type, подключенных к этому клиенту.

name_based_restore(...): Восстанавливает из checkpoint_prefix DTensors, основанные на имени.

name_based_save(...): Сохраняет тензор, основанный на имени, в контрольной точке.

num_clients(...): Возвращает количество клиентов в кластере DTensor.

num_global_devices(...): Возвращает количество устройств типа device_type в кластере DTensor.

num_local_devices(...): Возвращает количество устройств типа device_type, подключенных к текущему клиенту.

pack(...): Упаковывает компоненты tf.Tensor в DTensor.

relayout(...): Изменяет расположение tensor.

run_on(...): Выполняет заключённые функции в области устройства DTensor.

sharded_save(...): Сохраняет фрагменты указанных тензоров с именами безопасным для многоклиентских операций способом.

shutdown_tpu_system(...): Выключает систему TPU.

unpack(...): Распаковывает DTensor в компоненты tf.Tensor.