tf.compat.v1.tpu.replicate

Создаёт оператор графа, который выполняет реплицированное вычисление на TPU.

tf.compat.v1.tpu.replicate(
    computation: Callable[..., Any],
    inputs: Optional[List[List[core_types.Tensor]]] = None,
    infeed_queue: Optional[tpu_feed.InfeedQueue] = None,
    device_assignment: Optional[tf.tpu.experimental.DeviceAssignment] = None,
    name: Optional[Text] = None,
    maximum_shapes: Optional[Any] = None,
    padding_spec: Optional[tf.compat.v1.tpu.PaddingSpec] = None,
    xla_options: Optional[tf.tpu.XLAOptions] = None
) -> List[Any]

Пример базового использования, где inputs имеет статическую форму:

def computation(x):
  x = x + 1
  return tf.math.reduce_mean(x)

x = tf.convert_to_tensor([1., 2., 3.])
y = tf.convert_to_tensor([4., 5., 6.])
tf.compat.v1.tpu.replicate(computation, inputs=[[x], [y]])

Если у inputs динамические формы, и вы хотите автоматически букетировать входные данные для избежания перекомпиляции XLA. См. пример ниже:

def computation(x):
  x = x + 1
  return tf.math.reduce_mean(x)

# Assume input tensors in two replicas `x` and `y` both have dynamic shape
# ([None, 2]).
tf.compat.v1.tpu.replicate(
  computation,
  inputs=[x, y],
  maximum_shapes=[tf.TensorShape([None, None])],
  padding_spec=tf.compat.v1.tpu.PaddingSpec.POWER_OF_TWO)

Аргументы
`computation`	Функция Python, которая строит вычисление для репликации.
`inputs`	Список списков входных тензоров или `None` (эквивалентно `[[]]`), индексированные по `[replica_num][input_num]`. Все реплики должны иметь одинаковое количество входов. Каждый вход может быть вложенной структурой, содержащей значения, преобразуемые в тензоры. Обратите внимание, что передача списка размерности N совместимых значений приведет к списку размерности N скалярных тензоров, а не к одному тензору ранга N. Если вам требуется другое поведение, преобразуйте часть входов в тензоры с помощью `tf.convert_to_tensor`.
`infeed_queue`	Если не `None`, кортеж аргументов, которые нужно добавить как входы в вычисление.
`device_assignment`	Если не `None`, описание сопоставления между логическими ядрами в вычислении и физическими ядрами в топологии TPU. Использует значение по умолчанию для назначения устройств, если `None`. `DeviceAssignment` можно опустить, если каждая реплика вычисления использует только одно ядро, и либо есть только одна реплика, либо число реплик равно числу ядер в системе TPU.
`name`	(Устарело) Ничего не делает.
`maximum_shapes`	Вложенная структура tf.TensorShape, представляющая форму, до которой должна быть дополнена соответствующая компонента каждого элемента входных данных в каждой реплике. Любые неизвестные измерения (например, tf.compat.v1.Dimension(None) в tf.TensorShape или -1 в объекте типа тензора) будут дополнены до максимального размера этого измерения по всем репликам. Структура `maximum_shapes` должна быть такой же, как и у `inputs[0]`.
`padding_spec`	Перечисление, задаваемое `tpu.PaddingSpec`. Это описывает политику заполнения, когда `inputs` от `tpu.replicate` динамическая. Одно из применений – включить автоматическую букетизацию входных данных, установив значение в `tpu.PaddingSpec.POWER_OF_TWO`, что может помочь уменьшить перекомпиляцию в XLA.
`xla_options`	Экземпляр `tpu.XLAOptions`, который указывает параметры, передаваемые компилятору XLA. Используйте `None` для параметров по умолчанию.

Возвращаемое значение

Возвращаемое значение
Список выходных данных, индексированных по `[replica_num]`. Каждый выход может быть вложенной структурой, такой же, как возвращает computation(), за исключением нескольких случаев. Исключения включают: 1) Выход None: возвращается NoOp, который зависит от вычисления. 2) Выход с одним значением: возвращается кортеж, содержащий значение. 3) Выходы, состоящие только из операций: возвращается NoOp, который зависит от вычисления.

Список выходных данных, индексированных по [replica_num]. Каждый выход может быть вложенной структурой, такой же, как возвращает computation(), за исключением нескольких случаев.

Исключения включают:

1) Выход None: возвращается NoOp, который зависит от вычисления. 2) Выход с одним значением: возвращается кортеж, содержащий значение. 3) Выходы, состоящие только из операций: возвращается NoOp, который зависит от вычисления.

Исключения
`ValueError`	Если у всех реплик неравное количество входных тензоров.
`ValueError`	Если количество входов на реплику не совпадает с количеством формальных параметров для `computation`.
`ValueError`	Если статические `inputs` измерения не совпадают со значениями, заданными в `maximum_shapes`.
`ValueError`	Если структура входов на реплику не совпадает со структурой `maximum_shapes`.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/compat/v1/tpu/replicate