// interface the `CallOpInterface` so it should be handled separately.
    if (auto batch_function = dyn_cast<TF::BatchFunctionOp>(op)) {
      // Propagate the analysis results from within the callee's body.
      PropagatePotentiallyWrittenUpFromCallee(batch_function.func().getRegion(),
                                              batch_function.getOperands());
      return;
    }

  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
  // 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]

}

def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";

  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a bytecode function.

    Example:
      %res = tf_mlrt.batch_function(%input, %captured_input)  {
          device = "/device:CPU:0",

}

// -----

// Test lowering BatchFunction op.

func.func @batched_function(%x: tensor<?xi32>) -> tensor<?xi32> {
  return %x : tensor<?xi32>
}

// CHECK-LABEL: func @main
func.func @main(%x: tensor<1xi32>) -> (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) {
  // CHECK: [[y_future:%.*]] = tf_mlrt.batch_function
  // CHECK-SAME: f = @batched_function
  // CHECK-SAME: \22batch_function\22

}

def BatchFunctionOp : FallbackAsync_Op<"batch_function",
    [Pure, CoreRT_TypedAttributeTrait]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a BEF function.

    Example:
      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {

// CHECK-LABEL: func private @batched_function
func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %2 : tensor<1x3xf32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
  // CHECK:  "tf.BatchFunction"
  // CHECK-SAME: allowed_batch_sizes = [6]

  CopyDeviceAndUnderscoredAttributes(case_op, case_region);

  for (const auto& item : llvm::enumerate(case_region.getBranches())) {
    CreateCall(case_op, case_op.branch_function(item.index()),
               /*caller_region=*/item.value(), case_op.getInput(),
               /*use_region_args=*/false,
               /*forward_block_args=*/false);
  }

  # 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],

      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, flat_sym_ref);
    }
    // TODO(b/204997177): Deprecate and remove.
    func::FuncOp branch_function(int index) { return ResolveBranchFunction(nullptr, index); }

    // Resolve all branch functions.
    // Prefer passing in SymbolTableCollection to reduce lookup costs by