// Device types.
inline constexpr llvm::StringRef kTpuDevice = "TPU";
// _xla_outside_compilation
inline constexpr llvm::StringRef kXlaOutsideCompilationAttr =
    "_xla_outside_compilation";
// device attr
inline constexpr llvm::StringRef kDeviceAttr = "device";
// Function attribute to signal that a function should be skipped from TPU
// island outlining. The attribute is set in

// The main function's name in the serialized stablehlo module embedded in
// XlaCallModule's `module` attribute.
constexpr llvm::StringRef kStablehloMainFunctionName = "main";

// After deserializing the stablehlo functions from XlaCallModule,
// this XlaCallModule attribute refers to the deserialized stablehlo main
// function.
constexpr llvm::StringRef kStablehloEntryFunctionAttrName = "_entry_function";

  static constexpr size_t kResultIndex = 0;

  // Shapes of the input arguments.
  static constexpr int32_t kArg0Shapes[] = {
1, 2
  };
  static constexpr int32_t kArg1Shapes[] = {
3, 4
  };
  static constexpr int32_t kArg2Shapes[] = {
1
  };
  static constexpr int32_t kArg3Shapes[] = {
1
  };
  static constexpr int32_t kArg4Shapes[] = {
5
  };

using tpu::FunctionToHloArgs;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

static constexpr char kCompilationTimeStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/phase2_compilation_time";
static constexpr char kFullBridge[] = "full_bridge";
static constexpr char kCompilationStatusStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";

using ::tensorflow::monitoring::testing::CellReader;

static constexpr char kFailedLegalizationStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_failed_legalization_op_count";
static constexpr char kNonStaticOpStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_count";
static constexpr char kNonStaticOpSkipStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_skip_count";

// the SessionInitializerOp.
inline constexpr StringRef kTfSavedModelInitializerTypeAttr =
    "tf_saved_model.initializer_type";

// Indicates that the initializer corresponds to the restore op.
inline constexpr StringRef kTfSavedModelInitializerRestoreType = "restore_op";

// Indicates that the initializer corresponds to the init op.
inline constexpr StringRef kTfSavedModelInitializerInitType = "init_op";

namespace mlir {
namespace TFL {
namespace tac {
namespace {
// CPU
constexpr float kCPUArithmeticUnitCost = 1.0;

// This basically assumes pure load/store. This is just fake data.
constexpr float kCPUCopyUnitCost = 0.5;

// Default values.
constexpr float kCPUDefaultFixedValuedCost = 10000.0;

// Quantized inference cost efficiency.

using ::testing::IsEmpty;
using ::testing::IsNull;
using ::testing::NotNull;
using ::testing::Optional;
using ::tsl::testing::StatusIs;

using AttrsAndConstraintsTest = ::mlir::quant::QuantizationTestBase;

constexpr absl::string_view kModuleStatic = R"mlir(
  module {
    func.func @main(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {

// with existing function's name.
inline constexpr StringRef kTfQuantSaveFuncName = "tf_quant__save";

// Name of the TensorFlow Operation to be fetched to save the variables to
// checkpoint. This save op follows the SavedModel's load semantics, so it
// should return the file prefix of the checkpoint as a string tensor.
inline constexpr StringRef kTfQuantSaveOpName = "tf_quant__save_op";

 private:
  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

TEST_F(VerifyNoOutsideCompilationMarkersPassTest, PassesValidOps) {
  static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<1xi32> {