namespace tensorflow {

using ::testing::_;
using ::testing::NiceMock;
using ::testing::Return;
using ::testing::Test;

constexpr char kOk[] = "OK";
constexpr char kInvalidArgument[] = "INVALID_ARGUMENT";
constexpr char kSuccess[] = "kSuccess";
constexpr char kFailure[] = "kFailure";

class MockMlirOptimizationPass : public MlirOptimizationPass {
 public:

// will be evicted on the next read.
constexpr char kMaxStaleness[] = "GCS_READ_CACHE_MAX_STALENESS";
constexpr uint64_t kDefaultMaxStaleness = 0;

constexpr char kStatCacheMaxAge[] = "GCS_STAT_CACHE_MAX_AGE";
constexpr uint64_t kStatCacheDefaultMaxAge = 5;
// The environment variable that overrides the maximum number of entries in the
// Stat cache.
constexpr char kStatCacheMaxEntries[] = "GCS_STAT_CACHE_MAX_ENTRIES";

      GetMlirModuleFromString(module_with_two_clusters, &context));
  auto clusters = GetClusters(module.get());
  constexpr StringRef kTarget0 = "tpu0";
  EXPECT_EQ(clusters.count(kTarget0), 1);
  EXPECT_EQ(clusters[kTarget0].size(), 1);
  EXPECT_EQ(clusters[kTarget0][0].ops.size(), 2);

  constexpr StringRef kTarget1 = "tpu1";
  EXPECT_EQ(clusters.count(kTarget1), 1);
  EXPECT_EQ(clusters[kTarget1].size(), 1);

#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace {

using ::mlir::OpPassManager;
using ::tensorflow::monitoring::testing::CellReader;
using ::testing::HasSubstr;

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

#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_
#define TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_

namespace tensorflow {
namespace tfrt_compiler {

// Use __ prefix to indicate this is internal attribute.
inline constexpr char kOpKeyAttrName[] = "__op_key";

}  // namespace tfrt_compiler
}  // namespace tensorflow

namespace tensorflow {

llvm::SmallVector<int64_t> ConvertTFShapeToMlir(llvm::ArrayRef<int64_t> shapes);

llvm::SmallVector<int64_t> ConvertMlirShapeToTF(llvm::ArrayRef<int64_t> shape);

static constexpr int64_t kTFDynamicSize = -1;
mlir::RankedTensorType GetTypeFromTFTensorShape(llvm::ArrayRef<int64_t> shape,
                                                mlir::Type elementType,

#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/tfrt/common/global_state.h"
#include "tensorflow/core/util/determinism.h"

namespace tensorflow {
namespace {
constexpr const char* kPjRtDeviceCompilerResourceName = "pjrt_device_compiler";
constexpr const char* kPjRtDeviceCompilationProfilerResourceName =
    "pjrt_device_compilation_profiler";
}  // namespace

absl::StatusOr<std::unique_ptr<Graph>> CreateSingleOpGraph(

namespace {

using ::stablehlo::quantization::DebuggerConfig;
using DebuggerType = DebuggerConfig::DebuggerType;

constexpr StringRef kEntryFuncAttrName = "_entry_function";
constexpr StringRef kOriginalEntryFuncAttrName = "_original_entry_function";
constexpr StringRef kCompositeFuncPrefix = "composite_";
constexpr StringRef kEmptyNodeName = "_empty_node";

// Returns a pair: `func_name` and `node_name` for the lifted function. In TF

using ::tsl::testing::IsOk;
using ::tsl::testing::StatusIs;

using SaveQuantizationReportInstrumentationTest = QuantizationTestBase;

TEST_F(SaveQuantizationReportInstrumentationTest, SaveReport) {
  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
      %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>

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

namespace mlir {
namespace TF {

namespace {

constexpr const char *kDeviceAttr = "device";
constexpr const char *kTFDeviceAttr = "tf.device";

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