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;

  auto resource_operand = assign_var_op.getOperand(0);
  auto var_handle_op =
      llvm::dyn_cast<mlir::TF::VarHandleOp>(resource_operand.getDefiningOp());
  if (!var_handle_op) {
    assign_var_op->emitRemark(
        "Operand idx 0 is not a tf.VarHandleOp. The initializing tensor is not "
        "saved to checkpoint.");
    return "";
  }

}

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

    }
  }
  return nullptr;
}

Type GetGlobalType(Operation* source) {
  if (auto var_handle_op = dyn_cast<TF::VarHandleOp>(source)) {
    // Resources are represented as tensor<resource<tensor<...>>>, so
    // unwrap until we get to the inner tensor<...>.
    auto tensor =
        llvm::dyn_cast<TensorType>(var_handle_op.getResource().getType());
    if (!tensor) return nullptr;
    TF::ResourceType resource =

    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 =

namespace {
Node* MakeRead(const Scope& scope, const string& id,
               Node** var_handle_op = nullptr) {
  Output var_handle =
      ops::VarHandleOp(scope.WithOpName("Var" + id), DT_FLOAT, TensorShape({}));
  Output read =
      ops::ReadVariableOp(scope.WithOpName("Read" + id), var_handle, DT_FLOAT);
  if (var_handle_op) {
    *var_handle_op = var_handle.node();
  }
  return read.node();
}

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