# RUN: tf-mlir-translate -graphdef-to-mlir -tf-upgrade-legacy %s -tf-output-arrays=hash_table_node,variable_node,variable_v2_node -o - | FileCheck %s

node: {
  name: "hash_table_node"
  op: "HashTableV2"
  attr: {
    key: "key_dtype"
    value: {
      type: DT_INT32
    }
  }
  attr: {
    key: "shared_name"
    value: {
      s: ""
    }
  }
  attr: {
    key: "value_dtype"
    value: {
      type: DT_FLOAT
    }

    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
node {
  name: "variable_node"
  op: "Const"
  input: "^input"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {
  %result = "tf.IfrtCall"(%arg0, %arg1) <{program_id = 1234 : i64, variable_arg_indices = [0 : i32, 1 : i32], variable_names = ["a", "b"]}> : (tensor<?xf32>, tensor<?xf32>) -> (tensor<1x1xf32>)
  func.return
}

// -----
func.func @test_ifrt_call_fail_unsorted_variable_arg_indices(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {

                        "Null Session provided.");
  std::vector<tensorflow::Tensor> resource_tensors;
  auto variable_names = GetVariableNames(var_handle_ops);
  if (variable_names.empty()) return resource_tensors;

  auto status = session->Run({}, variable_names, {}, &resource_tensors);
  if (!status.ok())
    return absl::Status(absl::StatusCode::kInternal, status.message());
  return resource_tensors;
}

    return failure();
  }

  // Fetch all VarHandleOp.
  llvm::StringSet<> variable_names;
  llvm::SmallVector<TF::VarHandleOp, 4> var_ops;
  for (auto func_op : module.getOps<func::FuncOp>()) {
    for (auto var_handle_op : func_op.getOps<TF::VarHandleOp>()) {
      auto variable_name = GetVariableName(var_handle_op);
      if (variable_names.count(variable_name)) continue;
      var_ops.emplace_back(var_handle_op);

  // Read all resource variables from the session.
  std::vector<std::string> variable_names;
  variable_names.reserve(resource_names.size());
  for (StringRef name : resource_names) variable_names.push_back(name.str());

  std::vector<Tensor> resource_tensors;
  Status status = session->Run(
      /*inputs=*/{}, variable_names,
      /*target_tensor_names=*/{}, &resource_tensors);
  if (!status.ok()) {