import tensorflow as tf

@tf.function
def simple_fn(x, y):
    return x + y

input_tensor = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32)
captured_tensor = tf.constant([4.0], dtype=tf.float32)

def wrapped_fn(in_tensors, captured_tensors):
    return simple_fn(in_tensors[0], captured_tensors[0])

defun_func = tf.function(wrapped_fn).get_concrete_function([input_tensor], [captured_tensor])


result = tf.raw_ops.BatchFunction(
    in_tensors=[input_tensor],
    captured_tensors=[captured_tensor],
    f=defun_func,
    num_batch_threads=tf.constant(93389718, dtype=tf.int32),
    max_batch_size=10,
    batch_timeout_micros=1000,
    Tout=[tf.float32],
    max_enqueued_batches=10,
    allowed_batch_sizes=[5, 10],
    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
)

print(result)

2024-06-13 14:42:13.314670: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-13 14:42:19.076124: F external/local_tsl/tsl/platform/default/env.cc:74] Check failed: ret == 0 (11 vs. 0)Thread batch_threads_ creation via pthread_create() failed.
Aborted (core dumped)