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

#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(

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

# Test importing a saved model with 2 signatures that are using a same
# BatchFunction Op, which references to a same inference_func from graph_def
# library. The result should be that both signatures uses the same
# BatchFunction Op (the shared_name is the same) and the same copy of
# inference_func.

# CHECK: f = @inference_func[[post_fix:[^,]*]]

// assigned inside the function called by "tf.BatchFunction".

module {
// CHECK-LABEL: @serving_default
  func.func @serving_default() -> (tensor<*xi32>) {
    // expected-remark@below {{device: "", container: "", shared_name: "var_0", is_potentially_written: false}}
    %0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xi32>>>
    %1 = "tf.BatchFunction"(%0) {
        f = @called_by_batch_func,

// the functions invoked by BatchFunction with the same shared_name and the
// function `compute_2` will not be removed as the shared_name is different.

// CHECK-LABEL: func private @batch_0
// CHECK: f = @compute_0
func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

# CHECK: "tf.BatchFunction"
# CHECK-SAME: shared_name = "batch_node"

# CHECK: func private @create_resource
# CHECK: tf.HashTableV2
# CHECK-SAME: shared_name = "hash_table_node@create_resource"
# CHECK: "tf.Variable"
# CHECK-SAME: shared_name = "variable_node"
# CHECK: "tf.VariableV2"
# CHECK-SAME: shared_name = "variable_v2_node"
# CHECK: "tf.BatchFunction"