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();

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 =

// GraphDef.
std::unique_ptr<OperationPass<ModuleOp>> CreateMergeSaveFunctionOpsToMainPass();

// Creates a pass that "unfreezes" ConstOps into variables. Each ConstOp's use
// will be replaced by a VarHandleOp -> ReadVariableOp pattern. The newly
// created variables will be initialized in the session initializer function via
// AssignVariableOps.
std::unique_ptr<OperationPass<ModuleOp>> CreateUnfreezeConstantsPass();

    resource variables directly..

    Here is a simple example in TensorFlow where a device doubles the value of a
    TensorFlow resource variable and returns new value:

    ```mlir
    %resource_handle = "tf.VarHandleOp"()
    %1 = "tf_device.cluster"() ( {
      %init_value = "tf.ReadVariableOp"(%resource_handle)
      "tf.AssignAddVariableOp"(%resource_handle, %init_value)
      %new_value = "tf.ReadVariableOp"(%resource_handle)

// Creates the per-device variables that represent the formatting state of each
// device.
llvm::SmallVector<TF::VarHandleOp, 4> CreateStateVars(
    const llvm::SmallDenseMap<llvm::StringRef, llvm::SmallVector<StringRef, 4>>&
        devices,
    Location loc, RankedTensorType key_type, OpBuilder* builder) {
  llvm::SmallVector<TF::VarHandleOp, 4> state_vars;

  // TODO(b/148913020): Remove this constraint once model parallelism is
  // supported.

// Test lowering of IfrtLoadVariableOp

// CHECK-LABEL: func @ifrt_load_variable_test
func.func @ifrt_load_variable_test() -> () {
  // CHECK: [[HANDLE:%.*]] = tf_mlrt.executeop()
  // CHECK-SAME:  VarHandleOp
  %0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
  // CHECK-NEXT: "tf_mlrt.ifrt_load_variable"([[HANDLE]])

  "tf_device.cluster"() ({
    // CHECK:      "tf.VarHandleOp"
    // CHECK-SAME: _xla_outside_compilation
    "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
    tf_device.return

  func.return %arg : tensor<!tf_type.resource<tensor<i32>>>
}

// CHECK-LABEL: func @side_effect_while_test
func.func @side_effect_while_test() -> (tensor<i32>) {
  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "c", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<i32>>>
  // CHECK: [[while_res:%.]]:2 = tfrt.while {{%.*}} @"side_effect_while_body_add2/tfrt_body_1"