tf.raw_ops.BatchFunction

Сгруппирует все входные тензоры для вычисления, выполняемого функцией.

tf.raw_ops.BatchFunction(
    in_tensors, captured_tensors, f, num_batch_threads, max_batch_size,
    batch_timeout_micros, Tout, max_enqueued_batches=10, allowed_batch_sizes=[],
    container='', shared_name='', batching_queue='',
    enable_large_batch_splitting=False, name=None
)

Например, в следующем коде

# This input will be captured.
y = tf.placeholder_with_default(1.0, shape=[])

@tf.Defun(tf.float32)
def computation(a):
  return tf.matmul(a, a) + y

b = gen_batch_ops.batch_function(
        f=computation
        in_tensors=[a],
        captured_tensors=computation.captured_inputs,
        Tout=[o.type for o in computation.definition.signature.output_arg],
        num_batch_threads=1,
        max_batch_size=10,
        batch_timeout_micros=100000,  # 100ms
        allowed_batch_sizes=[3, 10],
        batching_queue="")

Если несколько вызовов session.run одновременно пытаются вычислить b, значения a будут собраны, не определяемо конкатенированы по первой оси, и только один поток выполнит вычисление.

Предполагается, что все аргументы функции являются тензорами, которые будут сгруппированы по первому измерению.

Аргументы, которые захвачены, не группируются. Вызов session.run, выполняющий конкатенацию, будет использовать значения захваченных тензоров, доступные ему. Следовательно, типичные использования захваченных тензоров должны включать значения, которые остаются неизменными в вызовах session.run. Инференс является хорошим примером этого.

SparseTensor не поддерживается. Результат выполнения декорированной функции должен быть тензором или списком/кортежем тензоров.