OutsideCompilationClusterDependencies(
    const Graph* g, const string& outside_compilation_attr_name) {
  auto cluster_deps = std::make_unique<
      absl::flat_hash_map<string, absl::flat_hash_set<string>>>();

  for (const Edge* e : g->edges()) {
    auto src_outside_compilation =
        GetStringAttr(*e->src(), outside_compilation_attr_name);
    auto dst_outside_compilation =
        GetStringAttr(*e->dst(), outside_compilation_attr_name);

    // Create chain source and sink identity islands for current equivalence
    // class.
    auto chain_arg = func.getArgument(chain_index++);
    auto src_identity = builder_chain_src.create<TF::IdentityOp>(
        chain_arg.getLoc(), chain_arg.getType(), chain_arg);
    auto chain_src_island = CreateIsland(src_identity, {}, builder_chain_src);

    auto sink_identity = builder_chain_sink.create<TF::IdentityOp>(

  if (failed(GetLstmProperty(op, &lstm_variant, &lstm_property))) {
    return nullptr;
  }

  auto spec = std::make_unique<quant::OpQuantSpec>();

  for (const auto& enumerated_inputs : lstm_property.inputs) {
    int index = enumerated_inputs.first;
    auto& tensor_property = enumerated_inputs.second;
    if (tensor_property.use_derived_scale) {
      double scale = 1.0;

    // argument in the list.
    auto arg = bb.getArgument(0);

    auto remove_quantize_op = [&](QuantizeOp quantize_op) {
      auto quantize_output = quantize_op.getOutput();
      auto quantize_type = quantize_output.getType();
      input_types.push_back(quantize_type);
      auto new_arg = bb.addArgument(quantize_type, loc);
      quantize_output.replaceAllUsesWith(new_arg);

          auto dq_arg_type = dq_op.getArg().getType().cast<TensorType>();
          auto qtype = dq_arg_type.getElementType().cast<QuantizedType>();
          auto scast_op = rewriter.create<quantfork::StorageCastOp>(
              dq_op->getLoc(), dq_arg_type.clone(qtype.getStorageType()),
              dq_op.getArg());
          inputs.push_back(scast_op.getResult());
        } else if (!elem_type.isF32()) {

    llvm::SmallVector<Value, 4> new_operands;
    for (auto arg : llvm::enumerate(compose_func_type.getInputs())) {
      if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(arg.value())) {
        auto casted = builder.create<CastOp>(op->getLoc(), tensor_type,
                                             op->getOperand(arg.index()));
        new_operands.push_back(casted);
      } else if (auto list_type =

}

// Checks if an operation of the tf_executor dialect can have TPU devices
// propagated through.
bool IsSupportedExecutorOp(Operation& op) {
  auto ops_have_same_device = [](Operation* lhs, Operation* rhs) {
    auto lhs_device_attr = lhs->getAttrOfType<StringAttr>(kDeviceAttr);
    auto rhs_device_attr = rhs->getAttrOfType<StringAttr>(kDeviceAttr);
    return (!lhs_device_attr && !rhs_device_attr) ||

    llvm::SmallVectorImpl<std::string>* var_handle_shared_names) {
  Block& block = function.front();
  auto func_type = function.getFunctionType();

  auto func_arg_types = llvm::to_vector<4>(func_type.getInputs());
  llvm::SmallDenseMap<llvm::StringRef, int> var_arg_index_by_name;
  for (auto var_handle_op :
       llvm::make_early_inc_range(block.getOps<TF::VarHandleOp>())) {

    return signalPassFailure();

  // Eliminate TPUCompilationResultOp now that the rewrite is complete.
  for (auto& it : compilation_results) {
    for (auto op : it.second) {
      if (!op.use_empty()) {
        mlir::InFlightDiagnostic err = op.emitError("uses remain post rewrite");
        for (auto user : op->getUsers())
          err.attachNote(user->getLoc()) << "remaining user";
        return signalPassFailure();

namespace {

// Counter that tracks total number and `write_version` of SavedModels written.
auto* saved_model_write_counter = monitoring::Counter<1>::New(
    "/tensorflow/core/saved_model/write/count",
    "The number of SavedModels successfully written.", "write_version");

// Counter that tracks total number and `write_version` of SavedModels read.
auto* saved_model_read_counter = monitoring::Counter<1>::New(