TensorArraySideEffectAnalysis tensor_array_side_effect_analysis(module);

    for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
      func_op.emitRemark() << "HasAtMostTensorArrayEffect: "
                           << tensor_array_side_effect_analysis
                                  .HasAtMostTensorArrayEffect(func_op);
    }
  }
};

mlir::PassRegistration<TestTensorArraySideEffectAnalysis> pass;

};

void RemoveShardingCustomCallPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  populateWithGenerated(patterns);

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  if (failed(applyPatternsAndFoldGreedily(func_op, frozen_patterns))) {
    func_op.emitWarning() << "Failed to converge "

          ConvertToLegacyCompileAndReplicateAttributesPass> {
  void runOnOperation() override;
};

LogicalResult ConvertToLegacyAttributes(func::FuncOp func_op) {
  auto result = func_op->walk([&](mlir::Operation* op) {
    if (failed(TF::HasValidCompilationAndReplicationAttributes(*op)))
      return WalkResult::interrupt();
    if (op->hasAttr(TF::kReplicationInfoAttr)) {

void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op) {
  auto entry_func_attr =
      func_op->getAttrOfType<DictionaryAttr>(kTfEntryFunctionAttr);
  if (!entry_func_attr) return;

  auto entry_func_attrs = SmallVector<NamedAttribute>(entry_func_attr.begin(),
                                                      entry_func_attr.end());

  MLIRContext* ctx = func_op.getContext();
  for (auto& named_attr : entry_func_attrs) {

};

// Folds tf.IdentityOp and tf.IdentityNOp if op device and the argument devices
// from the defining ops match.
void TensorDeviceCopyConversionPass::runOnOperation() {
  func::FuncOp func_op = getOperation();

  auto should_fold_op_func = [&func_op](const Value &arg,
                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and

void ConvertNMSPaddedFunc::RewriteFunc() {
  func_->setAttr(kTFImplements,
                 StringAttr::get(func_.getContext(), kTfNMSPadded));
  Value boxes = func_.getArgument(0);
  Value scores = func_.getArgument(1);
  Value max_output_size = func_.getArgument(2);
  Value iou_threshold = func_.getArgument(3);
  Value score_threshold = func_.getArgument(4);
  auto output_type0 = func_.getFunctionType().getResult(0);

 public:
  explicit CostAnalysis(
      mlir::func::FuncOp func_op,
      const tfrt_stub::CostRecorder* cost_recorder = nullptr) {
    cost_recorder_ = cost_recorder;
    AnalyzeArguments(func_op);
    AnalyzeBlock(&func_op.front());
  }

  int64_t GetCost(mlir::Operation* op) const;

 private:
  void AnalyzeArguments(mlir::func::FuncOp func_op);
  void AnalyzeBlock(mlir::Block* block);

  op->removeAttr(kTPUEmbeddingAttr);
  return success();
}

LogicalResult FindTPUEmbeddingOps(
    func::FuncOp func_op, llvm::StringMap<Operation*>* enqueue_op_map,
    llvm::StringMap<Operation*>* recv_activation_op_map,
    llvm::StringMap<Operation*>* send_gradient_op_map) {
  auto walk_result = func_op.walk([&](Operation* op) {
    if (llvm::isa<TF::RecvTPUEmbeddingActivationsOp>(op))

  const StringSet<> noinline_functions = GetNoinlineFunctionsSet();

  func::FuncOp func_op = getOperation();
  Builder builder(&getContext());

  // Adds the `tf._noinline = true` attribute to the function if the name
  // matches.
  if (noinline_functions.contains(func_op.getSymName())) {
    func_op->setAttr(kTfNoinlineAttr, builder.getBoolAttr(true));
    LLVM_DEBUG(llvm::dbgs()

  auto func_op_quant_scale_spec = GetStableHloQuantConstraints;

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // The function might contain more stats ops than required, and it will
    // introduce requantize if the calibration stats have conflicts. This tries
    // to remove all the redundant stats ops.
    RemoveRedundantStatsOps(func_op, func_op_quant_spec,
                            func_op_quant_scale_spec);