namespace {

using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::mhlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;

static constexpr char kHasTpuPartitionedCallStreamzName[] =
    "/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call";

class InferenceMetricsPassTest : public ::testing::Test {
 protected:

  int64_t operands_size = get_size(inst->getOperandTypes());

  constexpr int kSizeFactor = 2;
// TODO(b/233827625): Remove TF_DISABLE_CONSTANT_FOLDING macro.
#ifdef TF_DISABLE_CONSTANT_FOLDING
  constexpr int64_t kResultsSizeThreshold = 0;
#else
  constexpr int64_t kResultsSizeThreshold = (1 << 23);  // 1 MB
#endif
  constexpr int64_t kOperandsSizeThreshold = (1 << 30);  // 128 MB

  }
  return attr;
}

constexpr char kFunctionalIfNodeName[] = "If";
constexpr char kFunctionalCaseNodeName[] = "Case";
constexpr char kFunctionalWhileNodeName[] = "While";
constexpr char kCompilableFunctionName[] = "CompilableFn";
constexpr char kCompilableFunctionNodeName[] = "n_c";
constexpr char kUncompilableFunctionName[] = "UncompilableFn";
constexpr char kUncompilableFunctionNodeName[] = "n_c_uncompilable";

using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::util::MessageDifferencer;
using ::tsl::testing::IsOk;
using ::tsl::testing::StatusIs;

using LiftAsFunctionCallTest = QuantizationTestBase;

constexpr absl::string_view kModuleLifted = R"mlir(
  module {

  SavedObjectGraph value;
  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(string(text_proto),
                                                          &value));
  return value;
}

constexpr absl::string_view kSingleChildFoo = R"(
nodes {
  children {
    node_id: 1
    local_name: "foo"
  }
  user_object {
    identifier: "_generic_user_object"
    version {
      producer: 1

    "/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status",
    "Counts how often a successful export from TF Dialect to Executor Dialect "
    "is",
    "status");

constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";

namespace {

void AddTfDialectToExecutorPasses(OpPassManager &pm) {
  pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());

    "/tensorflow/core/tf2xla/api/v1/tf_dialect_to_executor_dialect_status",
    "Counts how often a successful export from TF Dialect to Executor Dialect "
    "is",
    "status");

constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";

namespace {

void AddTfDialectToExecutorPasses(OpPassManager &pm) {
  auto add_pass = [&](std::unique_ptr<Pass> pass) {

#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"

namespace mlir {
namespace TFL {

constexpr char kTFImplements[] = "tf._implements";
constexpr char kLstmCellSimple[] = "LSTMCellSimple";
constexpr char kLayerNormalizedLstmCellSimple[] =
    "LayerNormalizedLstmCellSimple";
constexpr char kCoupleInputForgetGates[] = "CoupleInputForgetGates";

#include "tsl/lib/core/status_test_util.h"

namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {

constexpr char kExportStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status";
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";

using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;

 public:
  using TFRBase::TFRBase;
  static constexpr StringLiteral name = "tfr.tensor";
  static std::string getTypeName() { return "TFRTensorType"; }
};

class TFRTensorListType : public detail::TFRTypeImpl<TFRTensorListType> {
 public:
  using TFRBase::TFRBase;
  static constexpr StringLiteral name = "tfr.tensor_list";
  static std::string getTypeName() { return "TFRTensorListType"; }