"""Test file to display the error message and verify it with FileCheck."""

# RUN: %p/concrete_function_error | FileCheck %s

import sys

from absl import app

import tensorflow.compat.v2 as tf
if hasattr(tf, 'enable_v2_behavior'):
  tf.enable_v2_behavior()


class TestGraphDebugInfo(object):
  """Test stack trace can be displayed."""

  def testConcreteFunctionDebugInfo(self):

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)) {
      op->setAttr(TF::kTpuReplicateAttr, op->getAttr(TF::kReplicationInfoAttr));
      op->removeAttr(TF::kReplicationInfoAttr);
      op->removeAttr(TF::kCompileDeviceTypeAttr);
    }

                          op->getDialect()->getNamespace().str() +
                          " not in tf functional dialect";
      op->emitError() << error;
      return WalkResult::interrupt();
    }

    if (op->hasAttr(mlir::TF::kXlaOutsideCompilationAttr)) {
      std::string error =
          "op has outside compilation attribute _xla_outside_compilation which "
          "is not allowed after clustering";
      op->emitError() << error;

  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);

"""Test file to display the error message and verify it with FileCheck."""

# RUN: %p/saved_model_error | FileCheck %s

import sys

from absl import app

import tensorflow.compat.v2 as tf
if hasattr(tf, 'enable_v2_behavior'):
  tf.enable_v2_behavior()


class TestModule(tf.Module):
  """The test model has unsupported op."""

  @tf.function(input_signature=[tf.TensorSpec(shape=[3, 3], dtype=tf.float32)])

  void runOnOperation() override;
};

void RestoreFunctionNameFromXlaCallModuleOp(TF::XlaCallModuleOp& call_op,
                                            SymbolTable& symbol_table) {
  if (!call_op->hasAttr(kOriginalStablehloEntryFunctionAttrName)) {
    return;
  }

  const auto original_function_name = call_op->getAttrOfType<StringAttr>(
      kOriginalStablehloEntryFunctionAttrName);

// Marks the op as no fallback.
def MarkNoFallback : NativeCodeCall<"SetNoFallbackAttr($_builder, $0)">;

def NoFallbackAttrNotSet : Constraint<CPred<
  "!$0.getDefiningOp()->hasAttr(kNoFallbackAttr)">>;

class FloatValueEquals<string val> : Constraint<CPred<
  "FloatValueEquals($0, " # val # ")">>;

class RankEquals<string rank> : Constraint<CPred<
  "RankEquals($0, " # rank # ")">>;

                      : info->DirectControlSuccessors(op);
    };
    // Some tests have recursive calls and other shenanigans, so allow
    // them to skip side effect analysis.
    if (!func->hasAttr("ignore_side_effects_for_testing")) {
      info =
          &getAnalysis<mlir::TF::SideEffectAnalysis>().GetAnalysisForFunc(func);
    }
    func->walk([&](Operation* function) {

    // '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");
    }