// CHECK: %[[VAR0:.*]] = "tf.VarHandleOp"()
    %0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>

    // CHECK: "tf.AssignVariableOp"(%[[VAR0]], %[[ASSET0]])
    "tf.AssignVariableOp"(%0, %arg0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()

    // CHECK: %[[VAR1:.*]] = "tf.VarHandleOp"()

    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>
    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<100x50xf32>>>) -> tensor<100x50xf32>

// CHECK-NOT:  IfrtLoadVariable
// CHECK:      "tf.VarHandleOp"
// CHECK-NEXT: "tf.AssignVariableOp"
// CHECK-NEXT: "tf.ReadVariableOp"
// CHECK-NEXT: "tf.StatefulPartitionedCall"
// CHECK-NEXT:  return 
//
module {
  func.func @serving_default() -> tensor<*xi32> {
    %cst = "tf.Const"() <{value = dense<"some_test.txt"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>

// Test that VarHandleOp is not marked as "potentially written" if it is not
// assigned inside the function called by "tf.BatchFunction".

module {
// CHECK-LABEL: @serving_default
  func.func @serving_default() -> (tensor<*xi32>) {
    // expected-remark@below {{device: "", container: "", shared_name: "var_0", is_potentially_written: false}}

    %0 = "tf.VarHandleOp"() {container = "", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<3xf32>>>
    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3xf32>>>) -> tensor<3xf32>
    %2 = "tf.VarHandleOp"() {container = "", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<?xf32>>>

  }

  StringRef getDescription() const final {
    return "Removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns "
           "from the initializer function of type 'restore_op'.";
  }

  void runOnOperation() override;
};

// Finds and removes the `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)`
// pattern. `tf.VarHandleOp` and `tf.Const` are removed unless they are used by
// other ops.

  // CHECK-NOT: "tf.VarHandleOp"
  // CHECK-NOT: "tf.ReadVariableOp"
  // CHECK: %[[CONST:.*]] = "tf.Const"()
  // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%arg1, %[[CONST]])
  // CHECK: %[[PACK:[0-9]*]] = "tf.Pack"(%[[CONST]], %[[ADD]])
  // CHECK: return %[[PACK]]
  %0 = "tf.Const"() {value = dense<4.200000e+01> : tensor<f32>} : () -> tensor<f32>

  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    // CHECK: "tf.VarHandleOp"
    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "invalid", shared_name = "invalid_var"} : () -> tensor<!tf_type.resource<tensor<100x50xf32>>>

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

    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
    return
  }
}

// -----
// single variable: VarHandleOp is before RestoreV2

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