for (auto& op : vars) RecursiveRemove(op, erase_list, dead_ops);

  // Erases the ops.
  for (auto op : erase_list) op->erase();
}

void RemoveVariablesInSessionInitializerPass::runOnOperation() {
  ModuleOp module_op = getOperation();

  for (auto init_func_op : GetInitializerFunctions(module_op)) {
    if (!init_func_op) return;

    if (init_func_op.getBlocks().size() != 1) {

    : public impl::ConstantOpDeviceAssignmentPassBase<
          ConstantOpDeviceAssignmentPass> {
  void runOnOperation() override;
};

void ConstantOpDeviceAssignmentPass::runOnOperation() {
  ModuleOp module = getOperation();

  module.walk([&](TF::ConstOp op) {
    // Keep the ConstOp if the op already have the device attribute.
    if (StringAttr device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {

#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/core_runtime/tensor_handle.h"  // from @tf_runtime

namespace tfrt {
class CoreRuntime;
}

namespace mlir {
class ModuleOp;
}

namespace tensorflow {

struct TfrtFunctionCompileOptions : public TfrtCompileOptions {
  // Currently only SavedModel API inference uses the tpu_fuse_ops option
  TfrtFunctionCompileOptions() {

      }
    }
    return WalkResult::advance();
  });

  return success();
}

void XlaBroadcast::runOnOperation() {
  FuncOp func = getOperation();
  mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
  if (!module) return signalPassFailure();
  func.walk([&](ClusterOp cluster) {
    if (auto replicate = cluster->getParentOfType<ReplicateOp>()) {

  XlaComputation computation = GetTestXlaComputation();

  TF_ASSERT_OK_AND_ASSIGN(TupleOp root_tuple,
                          ImportXlaComputationIntoModule(computation));

  ModuleOp parent_module =
      root_tuple.getOperation()->getParentOfType<ModuleOp>();
  EXPECT_EQ(parent_module, *module_);
}

TEST_F(Tf2XlaRewriterTest, FailsWithoutRootTuple) {
  TF_ASSERT_OK(CreateMlirModule());

using ::tensorflow::quantization::OpSet;

class LiftQuantizableSpotsAsFunctionsDRQPass
    : public PassWrapper<LiftQuantizableSpotsAsFunctionsDRQPass,
                         OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      LiftQuantizableSpotsAsFunctionsDRQPass)

  // Constructor used by the PassRegistration. This is only used by test.

  Object obj;
  obj.Set(
      String("_module"),
      Handle(impl::TaggedValue::Capsule(new mlir::OwningOpRef<mlir::ModuleOp>(
          std::move(module_or).value()))));

  auto get_string = [](Object self) {
    auto ref = self.Get<internal::Capsule>(String("_module"))
                   ->cast<mlir::OwningOpRef<mlir::ModuleOp>*>();
    return String(tensorflow::MlirModuleToString(ref->get(), false).c_str());
  };

// dump file.
template <typename FuncT>
absl::Status RunPasses(const absl::string_view name, FuncT add_passes_func,
                       mlir::MLIRContext& ctx, mlir::ModuleOp module_op) {
  mlir::PassManager pm{&ctx};
  add_passes_func(pm);

  mlir::StatusScopedDiagnosticHandler diagnostic_handler{&ctx};
  TF_RETURN_IF_ERROR(MaybeEnableIrPrinting(pm, name));

                                 "marker.";
      return failure();
    }
  }
  return success();
}

void XlaValidateInputsPass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symtab(module);
  llvm::SmallVector<func::FuncOp> entry_funcs = GetEntryFunctions(module);
  if (entry_funcs.empty()) {
    LOG(WARNING) << "missing entry functions";
  }

#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tsl/platform/statusor.h"

namespace mlir {
namespace tf2xla {
namespace internal {
namespace {

using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::mhlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;

static constexpr char kHasTpuPartitionedCallStreamzName[] =