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 не поддерживается. Возвращаемое значение декорированной функции должно быть тензором или списком/кортежем тензоров.