ImmediateTensorHandlePtr* handle) {
  ImmediateOpPtr varhandle_op(ctx->CreateOperation());

  TF_RETURN_IF_ERROR(varhandle_op->Reset("VarHandleOp", raw_device_name));
  TF_RETURN_IF_ERROR(varhandle_op->SetAttrType("dtype", dtype));

  // Note that if shape is unknown rank, shape.dim_sizes() will be empty, and
  // shape.dims() will be -1.

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

    %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"()

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

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

    // CHECK: tf_executor.island wraps "tf.VarHandleOp"() <{container = "a", shared_name = "x"}> {_mlir_name = "x", _output_shapes = [#tf_type.shape<>], allowed_devices = [], device = "/device:CPU:0", dtype = i64, shape = #tf_type.shape<>} : () -> tensor<!tf_type.resource<tensor<i64>>>

    %handle_0 = "tf.VarHandleOp"() {container="c", shared_name="a"} : () -> tensor<!tf_type.resource<tensor<1x10xf32>>>
    %handle_1 = "tf.VarHandleOp"() {container="c", shared_name="b"} : () -> tensor<!tf_type.resource<tensor<1x10xi64>>>

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