#include <string>

#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/core/graph/graph.h"

namespace tensorflow {
// The number of compiler threads to use for asynchronous device compilation.
inline constexpr int64_t kNumAsyncDeviceCompilerThreads = 10;

enum class DeviceCompileMode {
  kLazy,
  kStrict,
  kAsync,
};

enum class DeviceCompileState {
  kUncompiled,
  kCompiling,
  kCompiled,
};

// The Control type is a token-like value that models control dependencies
class ControlType : public Type::TypeBase<ControlType, Type, TypeStorage> {
 public:
  using Base::Base;
  static constexpr StringLiteral name = "tfl.control";
};

#include "tensorflow/compiler/mlir/lite/ir/tfl_ops_interface.h.inc"

}  // end namespace TFL
}  // end namespace mlir

#define GET_OP_CLASSES

namespace tensorflow {
namespace internal {
namespace {

using ::tensorflow::quantization::PyFunctionLibrary;
using ::tflite::optimize::ReducedPrecisionSupport;

// Op def string for TFLite_Detection_PostProcess Op.
constexpr mlir::StringRef kDetectionPostProcessOp =
    "name: 'TFLite_Detection_PostProcess' input_arg: { name: "
    "'raw_outputs/box_encodings' type: DT_FLOAT } input_arg: { name: "

#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

namespace mlir {
namespace tf_executor {

namespace {
constexpr llvm::StringRef kNestedModule = "_tpu_v1_compat_outlined";
constexpr llvm::StringRef kOutlinedFuncPrefix = "_tpu_v1_compat_outlined_func";

#define GEN_PASS_DEF_TPUBRIDGEEXECUTORISLANDOUTLININGPASS

// LINT.IfChange(random_access_file_ops_version)
constexpr int TF_RANDOM_ACCESS_FILE_OPS_API = 0;
constexpr int TF_RANDOM_ACCESS_FILE_OPS_ABI = 0;
constexpr size_t TF_RANDOM_ACCESS_FILE_OPS_SIZE =
    sizeof(TF_RandomAccessFileOps);
// LINT.ThenChange()

// LINT.IfChange(writable_file_ops_version)
constexpr int TF_WRITABLE_FILE_OPS_API = 0;
constexpr int TF_WRITABLE_FILE_OPS_ABI = 0;

using ::tsl::protobuf::TextFormat;

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Default platform for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kDefaultShardingValue[] = "";
constexpr char kMirroredVariableIndicesAttr[] = "_mirrored_variable_indices";

std::string GetRandomStateVariableName() {

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

constexpr StringRef kTfQuantSaveV2OpName = "tf_quant__save_save_v2";
constexpr StringRef kTfQuantSaveReturnOpName = "tf_quant__save_return";

// A pass that creates a new function that wraps the newly created SaveV2 op.

using mlir::stablehlo::ConstantOp;
using mlir::stablehlo::ConvertOp;
using ::stablehlo::quantization::QuantizationComponentSpec;

// Min/Max values used for creating ConstantOp.
constexpr float kMaxFloat16Value = 65504.f;
constexpr float kMinFloat16Value = -65504.f;

class QuantizeWeightPass
    : public impl::QuantizeWeightPassBase<QuantizeWeightPass> {
 public:
  explicit QuantizeWeightPass(

        AddUse(some_operation(rng), some_tensor(rng));
      }
    }
  };
};

TEST_F(RematSimulationTest, SimulationAgreesWithReality) {
  constexpr int kNumOperations = 128;
  constexpr int kNumTensors = 32;
  constexpr int kNumUses = kNumOperations * kNumTensors / 4;

  std::mt19937 rng;
  for (int i = 0; i < 1024; ++i) {
    RandomRemat remat(kNumOperations, kNumTensors, kNumUses, rng);