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 имеет значение true, выходные данные всех фрагментов 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/api_docs/python/tf/compat/v1/tpu/shard