}
  }
  // Update resource tensors.
  for (VarHandleOp var_handle_op : var_handle_ops) {
    builder.setInsertionPoint(var_handle_op);
    auto output_type = UnrankedTensorType::get(TF::ResourceType::get(
        {mlir::cast<TensorType>(ref_qtype)}, builder.getContext()));
    auto new_var_handle_op = builder.create<VarHandleOp>(
        var_handle_op.getLoc(), output_type, var_handle_op.getContainer(),
        var_handle_op.getSharedName());

    llvm::ArrayRef<TF::VarHandleOp> var_handle_ops) {
  std::vector<std::string> names;
  names.reserve(var_handle_ops.size());
  for (auto var_handle_op : var_handle_ops)
    names.push_back(GetVariableName(var_handle_op));
  return names;
}

tensorflow::Var* GetVariableFromSession(mlir::TF::VarHandleOp var_handle_op,
                                        llvm::StringRef device_name,

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

  // hoisted.
  auto read_var_op = llvm::dyn_cast<ReadVariableOp>(op);
  if (!read_var_op) return false;
  auto var_handle_op = llvm::dyn_cast_or_null<VarHandleOp>(
      read_var_op.getResource().getDefiningOp());
  if (!var_handle_op) return false;
  return read_only_vars.contains(GetResourceHandle(var_handle_op));
}

bool ShouldMoveOutOfRegion(
    Operation *op, Region *region,

    return failure();
  }

  // Fetch all VarHandleOp.
  llvm::StringSet<> variable_names;
  llvm::SmallVector<TF::VarHandleOp, 4> var_ops;
  for (auto func_op : module.getOps<func::FuncOp>()) {
    for (auto var_handle_op : func_op.getOps<TF::VarHandleOp>()) {
      auto variable_name = GetVariableName(var_handle_op);
      if (variable_names.count(variable_name)) continue;

}

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

    StringRef variable_name =
        GetNodeNameFromClassAttrOrSharedNameAttr(variable_v2_op);
    if (variable_name.empty()) {
      return signalPassFailure();
    }
    VarHandleOp var_handle_op = builder.create<VarHandleOp>(
        variable_v2_op.getLoc(),
        ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(ArrayRef<TensorType>{tensor_type},

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 =

  if (auto pop_back = dyn_cast<TF::TensorListPopBackOp>(op)) {
    return InferShapeForTensorListPopBackOp(pop_back);
  }

  if (auto var_handle_op = dyn_cast<VarHandleOp>(op)) {
    return InferShapeForVarHandleOp(var_handle_op);
  }

  if (auto xla_reduce_window_op = dyn_cast<XlaReduceWindowOp>(op)) {
    return InferShapeForXlaReduceWindowOp(xla_reduce_window_op);
  }