tf.compat.v1.tpu.replicate

Создаёт оператор графа, который выполняет реплицированное вычисление TPU.

tf.compat.v1.tpu.replicate(
    computation, inputs=None, infeed_queue=None, device_assignment=None, name=None,
    maximum_shapes=None, padding_spec=None, xla_options=None
)

Пример базового использования, в котором 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)