// CHECK: "tf.BatchFunction"(%arg0, %0)
  // CHECK: operandSegmentSizes = array<i32: 1, 1>

  %cst = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
  // CHECK:  "tf.BatchFunction"

  // CHECK: [[y_future:%.*]] = tf_mlrt.batch_function
  // CHECK-SAME: f = @batched_function
  // CHECK-SAME: \22batch_function\22
  %y = "tf.BatchFunction"(%x) {
    allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64,
    batching_queue = "", container = "", device = "/device:CPU:0",
    enable_large_batch_splitting = false, f = @batched_function,
    max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64,

  // CHECK-NEXT: return [[rendezvous_key_base]]
  %0 = "tf.BatchFunction"(%arg0, %arg0) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>

      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {
          num_batch_threads = 1,
          max_batch_size = 4,
          allowed_batch_sizes = [2 , 4],
          batch_timeout_micros = 50,
          container = "container",
          shared_name = "shared_name",
          batching_queue = "batching_queue",
          enable_large_batch_splitting = false,
          Tin = [f32],

     %arg105, %arg106, %arg107, %arg108, %arg109, %arg110, %arg111, %arg112) {
      num_batch_threads = 16,
      max_batch_size = 1,
      allowed_batch_sizes = [1],
      batch_timeout_micros = 0,
      container = "container",
      shared_name = "shared_name",
      batching_queue = "batching_queue",
      enable_large_batch_splitting = false,
      Tin = [i32, i32, i32, i32, i32],

    %handle = "tf.VarHandleOp"() {shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<0xf32>>>
    %0, %1 = "tf.BatchFunction"(%arg, %handle) {f = @f_batch_callee, operandSegmentSizes = array<i32: 1, 1>, batch_timeout_micros = 1000, max_batch_size = 8, num_batch_threads = 2} : (tensor<1xf32>, tensor<!tf_type.resource<tensor<0xf32>>>) -> (tensor<*xf32>, tensor<*xf32>)
    func.return %0, %1 : tensor<*xf32>, tensor<*xf32>
  }