variables;
  for (auto var_handle_op : var_handle_ops)
    variables.insert(GetResourceKey(var_handle_op));

  llvm::SmallVector<Operation*, 4> work_list;
  for (auto var_handle_op : session_init_func.getOps<TF::VarHandleOp>()) {
    if (variables.count(GetResourceKey(var_handle_op)))
      work_list.push_back(var_handle_op);
  }

      auto resource_type = var_handle_op.getResource().getType();
      func_arg_types.push_back(resource_type);
      var_handle_op.getResource().replaceAllUsesWith(
          block.addArgument(resource_type, var_handle_op.getLoc()));
    } else {
      var_handle_op.getResource().replaceAllUsesWith(
          block.getArgument(it.first->getSecond()));
    }
    var_handle_op.erase();
  }

}

using LocalVarOp = std::variant<TF::VarHandleOp, TF::MlirLocalVarOp>;

Value LocalVarOp_resource(LocalVarOp &op) {
  if (auto var_handle_op = std::get_if<TF::VarHandleOp>(&op)) {
    return var_handle_op->getResource();
  } else {
    return std::get<TF::MlirLocalVarOp>(op).getResource();
  }
}

void LocalVarOp_erase(LocalVarOp &op) {
  if (auto var_handle_op = std::get_if<TF::VarHandleOp>(&op)) {
    var_handle_op->erase();

    // Creates a VarHandleOp -> ReadVariableOp pair for each ConstOp.
    const auto resource_type = RankedTensorType::get(
        /*shape=*/{}, /*elementType=*/TF::ResourceType::get(
            /*subtypes=*/llvm::ArrayRef<TensorType>{const_op.getType()},
            builder.getContext()));
    auto var_handle_op =
        builder.create<TF::VarHandleOp>(const_op.getLoc(),

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 =

// CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}}
// CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]])

// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{{{.*value = dense<2.000000e\+00> : tensor<8xf32>.*}}}>
// CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"()  {{.*shared_name = "const_1".*}}
// CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[CST_1]])

// CHECK-SAME: %[[ARG_0:.*]]: tensor<!tf_type.string> {tf_saved_model.index_path = ["__tf_file_prefix"]}

// CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_0".*}} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
// CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_1".*}} : () -> tensor<!tf_type.resource<tensor<4xi32>>>

  attributes {tf_saved_model.exported_names = ["main"]} {
  // CHECK: tf.VarHandleOp
  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable1"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
  %1 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable2"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>

    // CHECK: "tf.VarHandleOp"
    func.return
  }
}

// -----

// Test mutation detection propagates across function calls.

module {
  func.func @f() -> tensor<0xf32> {
    // CHECK: "tf.VarHandleOp"

  attributes {tf_saved_model.exported_names = ["handles_ambiguous_var_handles"]} {
    // CHECK: VarHandleOp
    %0 = "tf.VarHandleOp"() {container = "", shared_name = "foo"} : () -> tensor<!tf_type.resource<tensor<10xf32>>>
    func.call @assign_resource(%0) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> ()
    %1 = "tf.VarHandleOp"() {container = "", shared_name = "foo"} : () -> tensor<!tf_type.resource<tensor<10xf32>>>