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

# 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 {
      }
    }
  }

        %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>

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

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>

  int num_retvals = 0;
  std::vector<AbstractTensorHandle*> dummy_outputs;
  return op_ptr->Execute(absl::MakeSpan(dummy_outputs), &num_retvals);
}

// Op: VarHandleOp()
// Summary:
//
// Description:

    %0 = tf_executor.graph {
      %outputs_4, %control_5 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "Variable"} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
      %outputs_10, %control_11 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "Variable_1"} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>

#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {
namespace {

// Implementation of varhandle that binds a VarHandleOp to an XlaResource of the
// same name. It is not safe to use this op in a JIT context.
class XlaAotOnlyVarHandleOp : public XlaOpKernel {
 public:
  explicit XlaAotOnlyVarHandleOp(OpKernelConstruction* c);

# CHECK: [[v1:%.*]] = "tf.Identity"([[v0]])
# CHECK: [[handle:%.*]] = "tf.VarHandleOp"
# CHECK-SAME: shared_name = [[shared_name:".*"]]
# CHECK: "tf.AssignVariableOp"([[handle]], [[v1]])

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME: tf_saved_model.exported_names = ["key"]
# CHECK: tf.VarHandleOp
# CHECK-SAME: shared_name = [[shared_name]]


def Test():

// CHECK-NOT: "tf_saved_model.session_initializer"
"tf_saved_model.session_initializer"() { initializers = [@func_init] } : () -> ()

// CHECK-LABEL: _tfrt_resource_init
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
// CHECK: tfrt_fallback_async.set_resource
// CHECK-SAME: {device = "/device:CPU:0", index = 0 : i64}


// CHECK-LABEL: func @init
// CHECK-SAME: {tfrt.cost_threshold = 1 : i64}