// ... then the code should not crash.
// CHECK-LABEL: Running 'test_batch_returns_multiple_refs'
func.func @test_batch_returns_multiple_refs() -> !tfrt.chain {
  %ch0 = tfrt.new.chain

  %0 = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>

  %1, %2 = tfrt_fallback_async.batch_function device("/device:CPU:0") @returns_multiple_refs (%0) {

    // 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]

      // CHECK: tfrt_fallback_async.batch_function device([[DEVICE:.*]]) @batched_func ([[BATCHED_FUNC_ARG:%.*]])
      // CHECK-SAME: Tcaptured = [!corert.resource]
      // CHECK-SAME: Tin = []
      // CHECK-SAME: Tout = [f32]

}

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",

func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
  func.return %arg0: tensor<i32>
}

// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>) -> (tensor<i32>) {
  // CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function
  // CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]

}

// -----

// 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)  {

#include "tfrt/tensor/tensor.h"  // from @tf_runtime

namespace tensorflow {
namespace {

// Creates a BEF file with a program that runs
// tfrt_fallback_async.batch_function with a empty function forwarding inputs or
// outputs.
std::pair<tfrt::BefBuffer, tfrt::RCReference<tfrt::BEFFile>> CreateBefFile(
    tfrt::HostContext* host) {
  std::string file_path = GetDataDependencyFilepath(

              !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {

  %res:112 = tfrt_fallback_async.batch_function device("/device:CPU:0") @f (%arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10,%arg11, %arg12, %arg13, %arg14,