fused_ln_lstm_func_ = createLstmCompositeFunc(builder_.get(), true, false);
  }

  void TearDown() override {
    fused_lstm_func_.erase();
    fused_lstm_func_cifg_.erase();
    fused_ln_lstm_func_.erase();
    builder_.reset();
  }

  func::FuncOp fused_lstm_func_;
  func::FuncOp fused_lstm_func_cifg_;
  func::FuncOp fused_ln_lstm_func_;

  builder->create<tf_device::ReturnOp>(old_terminator->getLoc(),
                                       merged_execute_launch.getResults());
  old_terminator->erase();

  // Remove the original TPUExecute op.
  execute_launch.erase();

  // Move all regions from old parallel_execute to new parallel_execute.
  for (auto region : llvm::zip(new_parallel_execute_op->getRegions(),

        !cond.getArgument(i).use_empty()) {
      explicitly_consumed_ids.push_back(i);
    }
  }
  // Empty consumed_element_ids implies none of results is used.
  if (explicitly_consumed_ids.empty()) {
    while_op.erase();
    return true;
  }
  // If every element is consumed, one can't reduce any operand.
  if (explicitly_consumed_ids.size() == n) {
    return false;
  }

  // Build the dependency graph.

    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();
  } else {
    std::get<TF::MlirLocalVarOp>(op).erase();
  }
}

std::optional<LocalVarOp> IsLocalVarOp(Operation &op) {
  if (TF::MlirLocalVarOp mlir_local_var_op =
          llvm::dyn_cast<TF::MlirLocalVarOp>(&op)) {

    block.eraseArgument(argument_index_to_remove);
    argument_types.erase(argument_types.begin() + argument_index_to_remove);
  }

  // Rewrite return if there are variable writes.
  const int return_operands_size = return_operands.size();
  if (return_operands_size > num_results_before) {
    builder.create<func::ReturnOp>(return_op.getLoc(), return_operands);
    return_op.erase();
  }

        if (cluster_name.has_value() &&
            cluster_name.value() != target_cluster_name) {
          tpu_ops.erase(iter);
          return true;
        }
        if (!cluster_name.has_value() &&
            !tpu_ops.count(wrapper.getOperation())) {
          tpu_ops.erase(iter);
          return true;
        }
      }
    }
  }
  return false;
}

    if (failed(UpdateReadUses(read_var, partitioned_input, partitioned_read,
                              partitioned_reads, is_packed)))
      return failure();
    read_var->erase();
    if (partitioned_input->use_empty()) partitioned_input->erase();
  }
  return mlir::TF::RemoveSingletonParallelExecuteOp(parallel_execute, &builder);
}

void TPUResourceReadsWritesPartitioningPass::runOnOperation() {

    builder->setInsertionPoint(terminator);
    auto yield = builder->create<tf_executor::YieldOp>(
        terminator->getLoc(), terminator->getOperands());
    terminator->erase();

    // Create new island for each region.
    builder->setInsertionPoint(island_op);
    auto execute_island = builder->create<tf_executor::IslandOp>(
        island_op.getLoc(), yield.getOperandTypes(),

        // for CPU. Note the 'tpu' value should be the same as `TPU` defined in
        // tensorflow/python/saved_model/tag_constants.py.
        if (quantization_options.has_debugger_config()) {
          tags.erase("tpu");
        }
        py_function_library.SaveExportedModel(
            dst_saved_model_path, *exported_model, src_saved_model_path, tags,
            signature_def_map);

        return absl::OkStatus();
      },

}

void Insert(const int item, std::vector<int>& items) {
  const auto [iter, found] = Find(item, items);
  if (!found) items.insert(iter, item);
}

void Erase(const int item, std::vector<int>& items) {
  const auto [iter, found] = Find(item, items);
  if (found) items.erase(iter);
}

}  // namespace

int Rematerializer::AddOperation(const bool is_stateful) {
  operations_.emplace_back();