}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
  // CHECK:  "tf.BatchFunction"
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64

  %result_1 = tfrt_fallback_async.batch_function device("/device:CPU:0") @matmul_cpu (%a1, %b) {
      num_batch_threads = 1,
      max_batch_size = 4,
      allowed_batch_sizes = [4],
      batch_timeout_micros = 1000000,
      container = "container",
      shared_name = "shared_name",
      batching_queue = "batching_queue",
      enable_large_batch_splitting = false,
      Tin = [i32],
      Tcaptured = [i32],

    %1 = "tf.BatchFunction"(%0) {
        f = @called_by_batch_func,
        operandSegmentSizes = array<i32: 1, 0>,
        batch_timeout_micros = 1000,
        max_batch_size = 8,
        num_batch_threads = 2
      } : (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<*xi32>
    return %1 : tensor<*xi32>
  }

func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

        type: [DT_INT64]
      }
    }
  }
  attr: {
    key: "Tout"
    value: {
      list: {
        type: [DT_INT64]
      }
    }
  }
  attr: {
    key: "batch_timeout_micros"
    value: {
      i: 100000
    }
  }
  attr: {
    key: "f"
    value: {
      func: {
        name: "inference_function"
      }
    }
  }
  attr: {

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

  // CHECK: tfrt_fallback_async.batch_function device("/device:CPU:0") @batched_function
  // CHECK-SAME: Tin = [f32]
  // CHECK-SAME: Tout = [f32]
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64

      // CHECK-SAME: Tcaptured = [!corert.resource]
      // CHECK-SAME: Tin = []
      // CHECK-SAME: Tout = [f32]
      %outputs_2, %control_2 = tf_executor.island wraps "tf.BatchFunction"(%outputs_1) {batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 0, 1>, shared_name = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>

func.func @serving_default(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.BatchFunction"(%arg0, %arg1) {f = @batched_func, num_batch_threads = 1 : i64, max_batch_size = 2 : i64, batch_timeout_micros = 10000 : i64, operandSegmentSizes = array<i32: 1, 1>} : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  return %0 : tensor<1xf32>
}
// The contents of `@serving_default` should have been inlined to `@batch_func`.