// Make sure that there are no more `tf.VarHandleOp`s and `tf.AssignVariableOp`s
// in this function.
// CHECK-NOT: "tf.VarHandleOp"
// CHECK-NOT: "tf.AssignVariableOp"

// Only the large constant is replaced with the `tf.VarHandleOp ->
// tf.ReadVariableOp` pattern and others remain as `tf.Const`s.
// CHECK: @serving_default

    // CHECK: tf_executor.island
    %island = tf_executor.island {
      // expected-remark@below {{Result #0, ID 0 : 0, 1, 2}}
      %vh0 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> !tf_res
      // expected-remark@below {{Result #0, ID 1 : 0, 1, 2}}
      %vh1 = "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> !tf_res
      // expected-remark@below {{Result #0, ID 2 : 0, 1, 2}}

  // CHECK-NOT: tf.VarHandleOp
  // CHECK-NOT: tf.ReadVariableOp
  // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
  // CHECK: [[r:%.*]] = "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: return [[r]]

    %outputs_7, %control_8 = tf_executor.island wraps "tf.Const"() {value = dense<"MatMul/b_0"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
    %outputs_9, %control_10 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "MatMul/b_0"} : () -> tensor<!tf_type.resource<tensor<20x4096xf32>>>

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

  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"}) {

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 =

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

             %arg3: !tf_res_md_f32 {tf.device = "/device:TPU:1"}) {

    %0 = "tf.Const"() {value = dense<100> : tensor<i32>} : () -> tensor<i32>
    // CHECK: %[[STATE0:.*]] = "tf.VarHandleOp"()
    // CHECK-SAME: device = "/device:TPU:0"
    // CHECK: %[[STATE1:.*]] = "tf.VarHandleOp"()
    // CHECK-SAME: device = "/device:TPU:1"
    // CHECK: %[[WHILE:.*]] = "tf.WhileRegion"(
    %1 = "tf.WhileRegion"(%0) ({
       // Condition region