patterns.add<ConvertTFXlaCallModuleOp>(&getContext(), module_op);
    patterns.add<RemoveCustomCallWithSharding>(&getContext());
    if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

std::unique_ptr<OperationPass<ModuleOp>>
CreateLegalizeTFXlaCallModuleToStablehloPass() {
  return std::make_unique<TFXlaCallModuleOpToStablehloPass>();
}

  for (auto func : module_op.getOps<func::FuncOp>()) {
    if (failed(applyPatternsAndFoldGreedily(func, frozen_patterns))) {
      func.emitError() << "quant-preprocess-op failed.";
      signalPassFailure();
    }
  }
}

}  // namespace

// Creates an instance of the TensorFlow dialect PreprocessOp
// pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePreprocessOpPass(

  if (use_tf2xla_fallback_) {
    tf2xla_fallback_device_type = device_type_;
  }
  if (failed(legalizeTF(op, legalize_chlo_, tf2xla_fallback_device_type,
                        prefer_tf2xla_))) {
    signalPassFailure();
  }
}

}  // end namespace

std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFPass(
    bool legalize_chlo, std::optional<StringRef> tf2xla_fallback_device_type,
    bool prefer_tf2xla) {

  RewritePatternSet patterns(context);

  patterns.add<InsertWeightParamPattern>(context);

  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    signalPassFailure();
  }
}

}  // namespace

          graph.emitError()
              << "Collect special Tpu ops failed: "
              << "Graph contains op with inconsistent cluster info\n";
          signalPassFailure();
          return WalkResult::interrupt();
        }
      }

      ExcludeUnqualifiedIdentityOp(cluster_to_tpu_ops_map);

      // Iterate until fixed point on the block, as it may contain multiple

  ModuleOp module_op = getOperation();

  func::FuncOp main_func_op = GetMainFunction(module_op);
  if (!main_func_op) {
    module_op.emitError("Main function op not found.");
    return signalPassFailure();
  }

  func::FuncOp save_func_op = GetSaveFuncOp(module_op);
  if (!save_func_op) return;

  MergeSaveFunctionOpsToMain(save_func_op, main_func_op);

  patterns.add<AddCustomAggregationOp>(ctx, calib_opts_);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-insert-custom-aggregation-ops failed.";
    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>
CreateInsertCustomAggregationOpsPass(const CalibrationOptions &calib_opts) {

  // semantics and then set token input node names attributes on the supported
  // ops.
  if (failed(RewriteCommunicationOps(module)) ||
      failed(SetTokenInputAttrs(module))) {
    signalPassFailure();
    return;
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreatePrepareTpuComputationForTfExportPass() {
  return std::make_unique<PrepareTpuComputationForTfExportPass>();

  populateWithGenerated(patterns);

  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-converting-tf-xla-op-to-tf-op failed.";
    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTfXlaOpToTfOpPass() {
  return std::make_unique<ConvertTfXlaOpToTfOpPass>();
}

  patterns.add<CommuteTransposeWithEwiseOps>(&getContext());
  patterns.add<CommuteBothInputsTransposedWithEwiseOps>(&getContext());
  if (failed(applyPatternsAndFoldGreedily(module, std::move(patterns)))) {
    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreatePushTransposeThroughEwisePass() {
  return std::make_unique<PushTransposeThroughEwisePass>();
}

}  // namespace TFL