tf.compat.v1.tpu.batch_parallel

Разбивает computation по измерению пакета для параллельного выполнения.

tf.compat.v1.tpu.batch_parallel(
    computation: Callable[..., Any],
    inputs: Optional[List[List[Optional[core_types.Tensor]]]] = None,
    num_shards: int = 1,
    infeed_queue: Optional[tpu_feed.InfeedQueue] = None,
    device_assignment: Optional[tf.tpu.experimental.DeviceAssignment] = None,
    name: Optional[Text] = None,
    xla_options: Optional[tf.tpu.XLAOptions] = None
)

Удобная обёртка вокруг shard().

inputs должен быть списком тензоров или None (эквивалентно пустому списку). Каждый входной элемент делится на num_shards частей по измерению с индексом 0, и вычисления выполняются параллельно для каждого фрагмента.

Тензоры передаются всем фрагментам, если они лексически захвачены computation. Например:

x = tf.constant(7) def computation(): return x + 3 ... = shard(computation, ...)

Выходные данные всех фрагментов объединяются обратно по измерению с индексом 0.

Входные и выходные данные вычисления должны быть тензорами ранга не ниже 1.

Аргументы
`computation`	Функция Python, создающая вычисление, которое применяется к каждому фрагменту входных данных.
`inputs`	Список входных тензоров или None (эквивалентно пустому списку). Размер измерения с индексом 0 каждого тензора должен делиться на `num_shards` без остатка.
`num_shards`	Количество фрагментов.
`infeed_queue`	Если не `None`, `InfeedQueue`, из которого добавляется кортеж аргументов в качестве входных данных для `computation`.
`device_assignment`	Если не `None`, `DeviceAssignment`, описывающий соответствие между логическими ядрами в вычислении и физическими ядрами в топологии TPU. Если `None`, используется назначение устройств по умолчанию. `DeviceAssignment` можно не указывать, если каждый фрагмент вычисления использует только одно ядро и при этом имеется либо только один фрагмент, либо количество фрагментов равно количеству ядер в системе TPU.
`name`	(Устарело) Ничего не делает.
`xla_options`	Экземпляр `tpu.XLAOptions`, указывающий параметры, передаваемые компилятору XLA. Для параметров по умолчанию используйте `None`.

Возвращает
Список выходных тензоров.

Вызывает исключение
`ValueError`	Если `num_shards <= 0`

© 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/batch_parallel