return function->emitError(
        absl::StrCat("failed to fetch device manager: ", status.message()));

  // Fetch all varHandleOp in the function.
  llvm::SmallVector<TF::VarHandleOp, 4> var_ops;
  for (auto var_handle_op : block.getOps<TF::VarHandleOp>())
    var_ops.emplace_back(var_handle_op);

  // Get resources from Session.
  auto resource_tensors_or = GetResourcesFromSession(var_ops, session);

  StringRef shared_name;
};

ResourceKey GetResourceKey(TF::VarHandleOp var_handle_op) {
  ResourceKey resource_key;

  if (auto attr = var_handle_op->getAttrOfType<StringAttr>("device")) {
    resource_key.device = attr.getValue();
  }

  resource_key.container = var_handle_op.getContainer();
  resource_key.shared_name = var_handle_op.getSharedName();

  return resource_key;
}

      return signalPassFailure();
    }

    auto var_handle_ops =
        init_func_op.getBlocks().front().getOps<VarHandleOp>();
    llvm::SmallVector<VarHandleOp, 4> init_vars(var_handle_ops.begin(),
                                                var_handle_ops.end());
    RemoveVariables(init_vars);
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>

    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 "";
  }

// CHECK: func.func private @tf_quant__save(%[[ARG:.*]]: tensor<!tf_type.string>)
// CHECK: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"()
// CHECK-SAME: {{.*shared_name = "var_0".*}}
// CHECK: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"()
// CHECK-SAME: {{.*shared_name = "var_1".*}}

// ReadVariableOps are inserted for each VarHandleOp to read the tensor values.

}

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

    }
  }
  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 that the locs are properly set to their shared_names.
// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}>
// CHECK-LOC-SAME: loc("var_0")
// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_1".*}}}>
// CHECK-LOC-SAME: loc("var_1")
}

// -----

// RestoreV2 op not created for `AssignVariableOp(VarHandleOp, Const)` patterns
// in the initializer function of "init_op" type.