namespace tensorflow {

using ::mlir::failure;
using ::mlir::LogicalResult;
using ::mlir::success;

namespace {
constexpr absl::string_view kPartitionedCall = "TPUPartitionedCall";

LogicalResult HasAttr(
    const Graph& graph, const FunctionLibraryDefinition* function_library,
    const std::function<bool(const Graph& graph)>& predicate) {
  if (predicate(graph)) {
    return success();
  }

  return {"tf.entry_function",
          tf_saved_model::kTfSavedModelInitializerTypeAttr};
}

bool IsEntryFunction(func::FuncOp func) {
  for (const auto &attr : GetEntryFunctionAttributeNames()) {
    if (func->hasAttr(attr)) {
      return true;
    }
  }
  return false;
}

llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module) {
  llvm::SmallVector<func::FuncOp> entry_funcs;

  // Copy all referenced StableHLO functions to the new module.
  WalkResult result = WalkReachableFunctions(
      entry_func,
      [&](func::FuncOp f) -> WalkResult {
        if (!f->hasAttr(kFromXlaCallModuleAttrName)) {
          return WalkResult::advance();
        }

        auto cloned = llvm::cast<func::FuncOp>(builder.clone(*f));
        cloned->removeAttr(kFromXlaCallModuleAttrName);

    // 'tf.entry_function' attribute set.
    llvm::SmallVector<func::FuncOp, 4> candidate_funcs;
    for (auto& entrypoint : entrypoints) {
      if (entrypoint.second->hasAttr("tf.entry_function")) {
        candidate_funcs.push_back(entrypoint.second);
      }
    }

    if (candidate_funcs.empty()) {
      return fail(module, "No entrypoints found");
    }

    }
  }
  return success();
}

LogicalResult HasTopLevelCompilationMarker(
    llvm::SmallVector<func::FuncOp> &entry_funcs) {
  for (auto &entry_func : entry_funcs) {
    if (entry_func->hasAttr(mlir::TF::kCompileDeviceTypeAttr)) {
      entry_func->emitError() << "TF2XLA MLIR Non-replicated Phase 1 Bridge "
                                 "does not support top-level compilation "
                                 "marker.";

//===----------------------------------------------------------------------===//

mlir::LogicalResult IfrtCallOp::verify() {
  auto func = getOperation()->getParentOfType<mlir::func::FuncOp>();
  if (func != nullptr && func->hasAttr("tfrt_ifrt_serving.program_id")) {
    return emitOpError() << "cannot be nested inside an IFRT program";
  }

  for (mlir::Value arg : getArgs()) {
    if (mlir::isa<mlir::TF::ResourceType>(

 private:
  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override {
    // Remove `_tpu_replicate` attributes on each operation first.
    if (op->hasAttr(tensorflow::kTPUReplicateAttr)) {
      op->removeAttr(tensorflow::kTPUReplicateAttr);
      return success();
    }

    // Remove TPU operations.
    if (isa<TF::TPUReplicateMetadataOp, TF::TPUCompilationResultOp,

    # performance by avoiding unnecessary cloning
    if cloned_types is None:
        cloned_types = _CLONED_TYPES_CACHE

    original_type = field.type_
    if is_dataclass(original_type) and hasattr(original_type, "__pydantic_model__"):
        original_type = original_type.__pydantic_model__
    use_type = original_type
    if lenient_issubclass(original_type, BaseModel):

          }
        }
      }
    });
  }
  // Remove `kSkipIslandOutlining` attributes.
  for (func::FuncOp func_op : getOperation().getOps<func::FuncOp>()) {
    if (func_op->hasAttr(mlir::TF::kSkipIslandOutlining)) {
      func_op->removeAttr(mlir::TF::kSkipIslandOutlining);
    }
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>