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 {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace TFL {

namespace {

constexpr char kNgrams[] = "tftext:Ngrams";
constexpr char kWhitespaceTokenizer[] = "tftext:WhitespaceTokenizer";
constexpr char kCustomSgnnProjection[] = "tftext:custom:SgnnProjection";
constexpr char kTFImplements[] = "tf._implements";

using mlir::TF::FuncAttr;
using mlir::TF::StringType;

namespace quant {
namespace {

using ::testing::NotNull;
using ::testing::SizeIs;

using ConstantFoldingTest = ::mlir::quant::QuantizationTestBase;

TEST_F(ConstantFoldingTest, FoldLargeConstant) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module {
      func.func @test_fold_constant() -> (tensor<1024x24x24x3xf32>) {
        %zp = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>

#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

constexpr char kTFAPIImplements[] = "tf.api_implements";
constexpr char kTFTextAPIPrefix[] = "tftext:";
constexpr char kCustomSSDPostprocessing[] = "TFLite_Detection_PostProcess";
constexpr char kTfNMSPadded[] = "non_max_suppression_padded_v2";
constexpr char kCustomMaxUnpooling[] = "addons:MaxUnpooling2D";
constexpr char kCustomDenseImageWarp[] = "addons:DenseImageWarp";

      dyn_cast_or_null<IntegerType>(qi8_type.getStorageType())));
}

using IsOpFullyQuantizedTest = QuantizationTestBase;

TEST_F(IsOpFullyQuantizedTest, TrueIfOpFullyQuantized) {
  constexpr absl::string_view kFullyQuantizedAdd = R"mlir(
    func.func @fully_quantized_add(%arg0: tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {

  EXPECT_TRUE(uq_requant_op);

  return uq_requant_op;
}

TEST_F(TfToUniformAttributeUtilsTest, UniformQuantizedAddOpAttributes) {
  TfToUniformAttributeUtilsTestPeer test_peer(&ctx_);

  constexpr absl::string_view kAddOpExpr =
      R"mlir(
      %0 = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656"> : tensor<1x3x2x2x!tf_type.qint32>} : () -> tensor<1x3x2x2x!tf_type.qint32>

namespace mlir {
namespace TFTPU {
namespace {

#define GEN_PASS_DEF_TPURESOURCEREADSWRITESPARTITIONINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";

struct TPUResourceReadsWritesPartitioningPass
    : public impl::TPUResourceReadsWritesPartitioningPassBase<

    {tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});

// There were no MLIR ops so the old bridge was called successfully.
constexpr char kOldBridgeNoMlirSuccess[] = "kOldBridgeNoMlirSuccess";
// There were no MLIR ops so the old bridge was called but it failed.
constexpr char kOldBridgeNoMlirFailure[] = "kOldBridgeNoMlirFailure";

namespace {

// Time the execution of kernels (in CPU cycles). Meant to be used as RAII.

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace odml {

static constexpr std::string_view kStablehloModuleDefaultEntryFuncName = "main";
static constexpr std::string_view kStablehloFuncNamePrefix = "XlaCallModule";
static constexpr char kShardingAttr[] = "mhlo.sharding";
static constexpr char kShardingName[] = "Sharding";

class RemoveCustomCallWithSharding

 public:
  // Number of input arguments for the compiled computation.
  static constexpr size_t kNumArgs = {{ARG_NUM}};

  static constexpr size_t kNumResults = {{RESULT_NUM}};

  // Number of variables for the compiled computation.
  static constexpr size_t kNumVariables = {{VARIABLE_NUM}};

  // Byte size of each argument buffer. There are kNumArgs entries.