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='',
    low_priority_max_batch_size=0,
    low_priority_batch_timeout_micros=0,
    low_priority_allowed_batch_sizes=[],
    low_priority_max_enqueued_batches=0,
    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="")