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)