// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
// CHECK-SAME:       <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)

  // CHECK-LABEL: func @main
  func.func @main() -> tensor<f32> {
    // CHECK-NEXT: %[[VAR:.*]] = "tf.VarHandleOp"
    // CHECK-NEXT: %[[READ_VAR:.*]] = "tf.ReadVariableOp"(%[[VAR]])
    // CHECK-NEXT: return %[[READ_VAR]]
    // CHECK-NOT: "tf.Cast"

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-output-arrays=func_call -o - | FileCheck %s

node {
  name: "x"
  op: "VarHandleOp"
  device: "/CPU:0"
  attr {
    key: "container"
    value {
      s: "a"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 130 : i32}} {
func.func @main() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "n"} : () -> tensor<!tf_type.resource<tensor<8xf32>>>

        %0 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
        "tf.AssignVariableOp"(%0, %cst) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()

        %cst_0 = "tf.Const"() {device = "", value = dense<[3, 4, 5, 6]> : tensor<4xi32>} : () -> tensor<4xi32>

#include "tensorflow/core/platform/errors.h"

using tensorflow::tracing::MaybeSetOpName;

namespace tensorflow {
namespace ops {

// Op: VarHandleOp()
// Summary: Creates a handle to a Variable resource.
//
// Description:
Status VarHandleOp(AbstractContext* ctx, AbstractTensorHandle** resource,
                   DataType dtype, const PartialTensorShape shape,

  // CHECK-NOT: tf.RestoreV2
  %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
  %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
  return
}


// CHECK-LABEL: func.func @used_restore_remains
func.func @used_restore_remains() {

func.func @main() -> tensor<*x!tf_type.resource> attributes {tf.entry_function = {inputs = "", outputs = "func_call"}} {
  %0 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.VarHandleOp"() {container = "a", device = "/CPU:0", dtype = i64, shape = "tfshape$", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i64>>> loc("x")

module attributes {tf_saved_model.semantics} {
  func.func private @tf_quant__save(%arg: tensor<!tf_type.string>) -> () {
    tf_executor.graph {
      %out, %ctl = tf_executor.island wraps "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
      %out_0, %ctl_0 = tf_executor.island wraps "tf.ReadVariableOp"(%out) : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<2xf32>

node {
  name: "G4SM2HGMB5"
  op: "VarHandleOp"
  device: "/job:localhost/replica:0/task:0/device:GPU:0"
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@G4SM2HGMB5/Read/ReadVariableOp"
      }
    }
  }
  attr {
    key: "container"
    value {
      s: ""
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {