namespace internal {
namespace {

using mlir::DialectRegistry;
using mlir::LogicalResult;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;
using mlir::PassManager;
using mlir::func::FuncOp;

std::string TestDataPath() {
  return tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/tf2xla/internal/testdata/");
}

class LoggingHooksTest : public ::testing::Test {
 public:

      : ConversionTarget(ctx), converter_(converter) {
    markUnknownOpDynamicallyLegal([this](Operation* op) {
      // The FuncOp type can contain types that the op's operand and result
      // types do not contain.
      if (auto func = dyn_cast<func::FuncOp>(op)) {
        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
      } else if (auto bitcast_convert_op =

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

        if (f == entry_func) {

    } else if (isa<IdentityNOp, IdentityOp>(op)) {
      value = op->getOperand(res_index);
    } else if (auto call = dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func = dyn_cast<func::FuncOp>(
          call.resolveCallable(&symbol_table_collection_));
      if (!func) break;
      // Check if the function being called has been analyzed. if not,
      // we cannot backtrack the value further.

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {
using mlir::Operation;
using mlir::OperationPass;
using mlir::WalkResult;
using mlir::func::FuncOp;

#define GEN_PASS_DEF_VERIFYINPUTDIALECTTOEXECUTORPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h.inc"

class VerifyInputDialectToExecutorPass

    func::FuncOp ResolveThenFunction(::mlir::SymbolTableCollection* table) {
      if (table)
        return table->lookupNearestSymbolFrom<func::FuncOp>(*this, getThenBranchAttr());
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
        *this, getThenBranchAttr());
    }
    // TODO(b/204997177): Deprecate and remove.

  ModuleOp module = getOperation();
  if (llvm::size(functions_) == 0) {
    functions_ = {"main"};
  }
  SymbolTable symbolTable(module);
  for (const std::string& f : functions_) {
    func::FuncOp func = symbolTable.lookup<func::FuncOp>(f);
    if (!func) continue;

    // This routine should only be called when control flow operations are still
    // represented with TF IfOp and WhileOp operations. In this case, there

namespace mlir {
namespace quant {

using QuantParamsEntry = QuantizationInfo::QuantParams;

namespace {
class ImportQuantStatsPass
    : public PassWrapper<ImportQuantStatsPass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ImportQuantStatsPass)

  explicit ImportQuantStatsPass(OperationToName op_to_name)
      : op_to_name_(op_to_name) {}

            mlir::func::FuncOp func =
                llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
            if (!func) continue;
            next_values_to_visit.push_back(
                func.getArgument(next_op_operand_num));
          } else if (auto while_op =
                         llvm::dyn_cast_or_null<TF::WhileOp>(next_op)) {
            func::FuncOp func = while_op.body_function();

  pm.addNestedPass<func::FuncOp>(
      tf_executor::CreateTFExecutorIslandCoarseningPass());
  pm.addPass(CreateTFShapeInferencePass());

  // Assign optimal data layout to layout sensitive operations and delete
  // redundant transposes from the IR.
  LayoutOptimizationPipelineOptions layout_optimization_options;
  CreateLayoutOptimizationPipeline(pm.nest<func::FuncOp>(),