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