// CHECK: func @init()
  // CHECK-NOT: tf.VarHandleOp
  // CHECK-NOT: tf.Const
  // CHECK-NOT: tf.AssignAddVariableOp
  // CHECK: return
  "tf_saved_model.session_initializer"() { initializers = [@init] } : () -> ()
  func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"], tf_saved_model.initializer_type = "restore_op"} {

    // CHECK-DAG: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp"
    // CHECK-DAG: [[MG_HANDLE:%.*]] = "tf.VarHandleOp"
    // CHECK-DAG: [[MS_HANDLE:%.*]] = "tf.VarHandleOp"
    // CHECK-DAG: [[MOM_HANDLE:%.*]] = "tf.VarHandleOp"
    %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<f32>>>

    %var_0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
    "tf.AssignVariableOp"(%var_0, %cst_0) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
    return
  }
  // Nothing has been removed.
  // CHECK: @init_func_init_op
  // CHECK-NEXT: "tf.Const"
  // CHECK-NEXT: "tf.VarHandleOp"
  // CHECK-NEXT: "tf.AssignVariableOp"

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