void TPURewritePass::runOnOperation() {
  TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(getOperation(), &devices)))
    return signalPassFailure();

  // Collect compilation results.
  llvm::DenseMap<Attribute, SmallVector<TF::TPUCompilationResultOp, 1>>
      compilation_results;
  auto result_init = getOperation().walk([&](TF::TPUCompilationResultOp op) {

 public:
  explicit AnalyzeVariablesPass() {}
  void runOnOperation() override;
};

void AnalyzeVariablesPass::runOnOperation() {
  auto* context = &getContext();
  auto module = getOperation();
  bool legalize_to_tfl = true;

  module.walk([&](Operation* op) {
    // Skip ops that are supported natively by TFLite.
    if (IsSupportedTFLiteResourceOp(op)) return WalkResult::advance();

  // Find all the executes first, since we will mutate the nodes around each
  // execute in the same tf_device.replicate op.
  llvm::SmallVector<tf_device::LaunchOp, 8> execute_launches;
  getOperation().walk([&](tf_device::LaunchOp op) {
    if (op.WrapsSingleOp() &&
        llvm::isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(
            op.GetBody().front()))
      execute_launches.push_back(op);

      new_op_mapper.map(result, new_result);
    }
  }

  call_op.erase();
}

void UnwrapXlaCallModuleOpPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  SymbolTable symbol_table(module_op);

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    Region& function_body = func_op.getBody();

    function_body.walk([&](TF::XlaCallModuleOp call_op) {

    return op.getOperation();
  }
  auto op = builder.create<tfl::ConstOp>(loc, value);
  op->setAttr("tfl.is_variable", builder.getUnitAttr());
  if (tensor.quantization && !tensor.quantization->min.empty()) {
    if (auto stats_op =
            ConvertMinMaxToStatsOp(tensor, builder, op.getResult())) {
      return stats_op;
    }
  }
  return op.getOperation();
}

          TPUHostComputationExpansionPass> {
  void runOnOperation() override;
};

void TPUHostComputationExpansionPass::runOnOperation() {
  OpBuilder builder(&getContext());
  getOperation().walk([&](tf_device::ClusterOp cluster) {
    ExpandHeadOutsideCompiledOps(cluster, &builder);
  });
}

}  // anonymous namespace

std::unique_ptr<OperationPass<func::FuncOp>>

  if (auto whileOp =
          llvm::dyn_cast<WhileRegionOp>(this->getOperation()->getParentOp())) {
    if (&whileOp.getCond() == this->getOperation()->getParentRegion()) {
      // cut off the boolean (the condition itself) at the start
      return MutableOperandRange(
          this->getOperation(), 1,
          this->getOperation()->getOperands().size() - 1);
    }

      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointAfter(op->getParentOfType<func::FuncOp>());
      cloned_func = llvm::dyn_cast_or_null<func::FuncOp>(
          rewriter.clone(*func.getOperation()));
      manager.insert(cloned_func);
      rewriter.setInsertionPointToStart(&cloned_func.getBody().front());
      auto result_type =
          RankedTensorType::get({3}, rewriter.getType<TF::StringType>());

          UnfoldLargeSplatConstantPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnfoldLargeSplatConstantPass)

  void runOnOperation() override {
    auto module = getOperation();

    mlir::OpBuilder op_builder(&module.getBodyRegion());
    module.walk([&](mlir::arith::ConstantOp const_op) {
      MaybeUnfoldLargeSplatConstant(&op_builder, const_op);
    });
  }

 private:

void CastBf16OpsToF32Pass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  auto module_op = getOperation();

  patterns.add<CastBf16OpsToF32>(ctx);
  populateWithGenerated(patterns);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {