tf.compat.v1.tpu.shard

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

tf.compat.v1.tpu.shard(
    computation: Callable[..., Any],
    inputs: Optional[List[core_types.Tensor]] = None,
    num_shards: int = 1,
    input_shard_axes: Optional[List[int]] = None,
    outputs_from_all_shards: Union[bool, List[bool]] = True,
    output_shard_axes: Optional[List[int]] = None,
    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
) -> List[core_types.Tensor]

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

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

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

в качестве входных данных.

Если outputs_from_all_shards истинно, выходы всех фрагментов computation конкатенируются обратно вдоль их output_shard_axes. В противном случае каждый вывод берется из произвольного фрагмента.

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

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

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

Исключения
`ValueError`	Если num_shards <= 0
`ValueError`	Если len(input_shard_axes) != len(inputs)
`ValueError`	Если len(output_shard_axes) != len(выходы из `computation`)

© 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/versions/r2.9/api_docs/python/tf/compat/v1/tpu/shard