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`, `InfeedQueue`, из которого к вычислениям добавляется кортеж аргументов в качестве входных данных.
`device_assignment`	Если значение не равно `None`, `DeviceAssignment`, описывающий соответствие между логическими ядрами вычислений и физическими ядрами в топологии 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, зависящий от computation по управлению. 2) Единственное выходное значение: будет возвращен кортеж, содержащий это значение. 3) Только операции в качестве выходных данных: будет возвращен NoOp, зависящий от computation по управлению.

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

Исключения:

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

Вызывает исключения
`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