// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> !tfrt.chain
func.func @assign_variable() {
  // CHECK: [[ch1:%.*]], %results = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/device:CPU:0") "tf.VarHandleOp"
  // CHECK-NEXT: [[ch2:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(1) cost({{.*}}) device("/device:CPU:0") "tf.AssignVariableOp"
  // CHECK-NEXT: [[out_ch:%.*]] = tfrt.merge.chains [[ch1]], [[ch2]]

  # Regression test for an issue where VarHandleOp wasn't being properly
  # imported into MLIR for "no-op" node expansion.
  def testVarHandleOp(self):
    x = constant_op.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

    # Note: we purposely make multiple calls to VarHandleOp to exercise the
    # cached kernal lookup path that was exhibiting the VarHandleOp import
    # issue.

// RUN: tf-opt %s -split-input-file -tfl-analyze-variables-pass --cse | FileCheck %s

// CHECK: module attributes {tfl._legalize_tfl_variables = true}
module {
  func.func @f() -> tensor<*xi32> {
    %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<*xi32>>>
    %2 = "tf.ReadVariableOp"(%0) {dtype = i32} : (tensor<*x!tf_type.resource<tensor<*xi32>>>) -> tensor<*xi32>
    func.return %2 : tensor<*xi32>
  }
}

// readonly.

// CHECK-LABEL: readvariableop_is_hoisted_if_readonly2
// CHECK:       [[CST_0:%.*]] = "tf.Const"
// CHECK:       [[VAR_1:%.*]] = "tf.VarHandleOp"
// CHECK-SAME:  "shared_name_var1"
// CHECK:       [[VAR_2:%.*]] = "tf.VarHandleOp"
// CHECK-SAME:  "shared_name_var2"
// CHECK:       [[CST_1:%.*]] = "tf.Const"
// CHECK:       [[VAR_VAL:%.*]] = "tf.ReadVariableOp"([[VAR_1]])

}

// Returns the name of the VarHandleOp that provides input (possibly indirectly)
// to node with node_name. A typical indirect chain of nodes (that can occur due
// to graph inlining) is the following: VarHandleOp -> Identity -> Identity ->
// ReadVariableOp. Calling the function on any of these nodes would return the
// name of the VarHandleOp.
StatusOr<string> GetVarHandleName(

  } else if (auto next = llvm::dyn_cast<TF::IteratorGetNextOp>(op)) {
    resource = next.getIterator();
  }

  if (llvm::dyn_cast_or_null<TF::VarHandleOp>(resource.getDefiningOp())) {
    return;  // We're already directly after a VarHandleOp.
  }

  const TF::ResourceDataflowState* state =
      solver.lookupState<TF::ResourceDataflowState>(resource);
  if (!state) {

// Test the basic cases where all uses of a ref variable can be converted.

// CHECK-LABEL: @init
func.func @init() {
  // CHECK-NOT: tf.VariableV2
  // CHECK-NOT: tf.Assign

  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
  // CHECK-SAME: shared_name = "x"
  // CHECK: "tf.AssignVariableOp"([[handle]], {{%.*}})
  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>

  func.return %3 : tensor<*xi32>
}

}

// -----

module attributes {tf_saved_model.semantics} {

// CHECK-LABEL: func private @reorder_execute_arg_defining_ops
// CHECK: tf.VarHandleOp
// CHECK-NEXT: tf.ReadVariableOp
// CHECK-NEXT: tf.TPUCompileMlirAndExecute
func.func private @reorder_execute_arg_defining_ops(%arg0: tensor<1x3xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {

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

TEST(ShapeInferenceTest, WhileLoopWithResource) {
  // Graph:
  // x = resource_variable_ops.var_handle_op(dtype=dtypes.float32, shape=[2, 3])
  // y = control_flow_ops.while_loop(lambda _: true, lambda x: x, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto x =
        ops::VarHandleOp(scope.WithOpName("x"), DT_FLOAT, TensorShape({2, 3}));
    auto enter =