mlir::tf_device::LaunchOp launch = CreateLaunchForBlock(
      builder, cluster, /*before=*/false, launch_block, host_device);

  auto operand_not_in_launch = [&](mlir::OpOperand& operand) {
    return !launch.getOperation()->isProperAncestor(operand.getOwner());
  };
  for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
                               launch.getResults()))

  }

  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                   FlatSymbolRefAttr::get(original_function_name));
}

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

  MLIRContext* ctx = module_op.getContext();
  OpBuilder builder(ctx);
  SymbolTable symbol_table(module_op);

  // TODO - b/298966126: Improve this logic if needed.

class VerifyClusteringPass
    : public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
 public:
  void runOnOperation() override;
};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +

  auto stats_op = b.create<quantfork::StatisticsOp>(
      b.getUnknownLoc(), res, layer_stats, axis_stats, axis);
  res.replaceAllUsesWith(stats_op);
  stats_op.getOperation()->replaceUsesOfWith(stats_op, res);
}

void ImportQuantStatsPass::ImportAsStatsOps(OpBuilder b, Operation *op,
                                            int index,

    // (x - mean) * scale / sqrt(variance + epsilon) + offset,
    // is then to compute (x * multiplier) + (offset - mean * multiplier).

    auto epsilon = materializeEpsilon(
        bn_op.getOperation(), bn_op.getEpsilonAttr(), fp_type,
        bn_op.getVariance(), variance_type, rewriter);
    if (!epsilon) {
      return failure();
    }

    // Compute multiplier = scale / sqrt(variance + epsilon)

  explicit ConvertFuncToBfloat16Pass() = default;

 private:
  void runOnOperation() override;
};

void ConvertFuncToBfloat16Pass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext* context = func_op.getContext();
  RewritePatternSet patterns(context);

  BFloat16TypeConverter converter;
  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,

// tf_device.parallel_execute
//===----------------------------------------------------------------------===//

LogicalResult ParallelExecuteOp::verify() {
  ParallelExecuteOp op = *this;
  const auto& regions = op.getOperation()->getRegions();
  if (regions.empty()) {
    return op.emitOpError() << "must have at least one region.";
  }

  int output_index = 0;
  for (const auto& region_and_index : llvm::enumerate(regions)) {

 public:
  void runOnOperation() override;
};

void InputMetricsLoweringPass::runOnOperation() {
  bool has_dynamic_op = false;
  Operation* func_op = getOperation();

  func_op->walk([&](Operation* op) {
    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return WalkResult::advance();

    if (mlir::mhlo::IsDynamicPadderOp(abstractOp->getTypeID())) {

      if (failed(layout)) return;

      op->setAttr("layout", layout.value());
    }
  }
};

void InfeedsOpsXlaAdjustLayout::runOnOperation() {
  getOperation().walk(runOnInfeedOp);
}

}  // anonymous namespace

std::unique_ptr<mlir::OperationPass<func::FuncOp>>
CreateInfeedsOpsXlaAdjustLayoutPass() {
  return std::make_unique<InfeedsOpsXlaAdjustLayout>();
}

    auto fail = [&](Operation* op, std::string message) {
      op->emitError(message);
      signalPassFailure();
    };

    DenseMap<StringRef, func::FuncOp> entrypoints;
    auto module = getOperation();
    module.walk([&](func::FuncOp op) {
      auto visibility = SymbolTable::getSymbolVisibility(op);
      if (visibility != SymbolTable::Visibility::Public) return;
      entrypoints[op.getSymName()] = op;
    });