#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project

namespace tensorflow {
namespace quantization {

inline constexpr absl::string_view kTfQuantConstantUnfreezingStepName =
    "tf_quant_constant_unfreezing";
inline constexpr absl::string_view kTfQuantInsertRestoreOpStepName =
    "tf_quant_insert_restore_op";

absl::Status UnfreezeConstantsAndSaveVariables(absl::string_view checkpoint_dir,

 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "mlrt.compiler.future";
};

// The MLIR type represents a C++ mlrt::Promise.
class PromiseType
    : public mlir::Type::TypeBase<PromiseType, mlir::Type, mlir::TypeStorage> {
 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "mlrt.compiler.promise";
};

// Presence of this attribute will mark the function as quantization target.
inline constexpr StringRef kFusedFunctionAttr = "tf_quant.composite_function";
// The keyword to detect if this is a `NullAttribute`.
inline constexpr StringRef kNullAttributeValue = "N/A";

// Prefixes attached to lifted functions.
constexpr StringRef kQuantizedFuncPrefix = "quantized_";
constexpr StringRef kCompositeFuncPrefix = "composite_";

using ::tsl::protobuf::io::ArrayInputStream;

// Empty (default) `QuantizationSpecs` proto.
constexpr absl::string_view kSpecsEmpty = R"pb(specs
                                               [])pb";

// Configure `QuantizationSpecs` to disable quantization for all dot_general
// quantizable units.
constexpr absl::string_view kSpecsDisableAllDotGeneral =
    R"pb(specs
         [ {

namespace mlir {
namespace TFL {
namespace tac {

// The copy can be non-consectutive copy. This is just fake data.
constexpr float kNNAPICopyUnitCost = 0.2;

// Default values.
constexpr float kNNAPIDefaultFixedValuedCost = 10000.0;

constexpr char NNAPIHardware::kId[];  // Define kId.

mlir::RewritePatternSet NNAPIHardware::GetTransformations(
    MLIRContext* context) const {

    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  Env* env_{};
  mlir::MLIRContext ctx_{};
};

TEST_F(SaveVariablesToCheckpointTest, VariableSavedToCheckpoint) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module attributes {tf_saved_model.semantics} {
      "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()

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

namespace mlir::quant {

constexpr char kAttrMapAttribute[] = "attr_map";

// Name of the string attribute attached to `XlaCallModuleOp`, which is the
// textproto representation of `Method`.
inline constexpr StringRef kQuantizationMethodAttr = "_quantization_method";

// Permutation from the NHWC tensor format to NCHW. This is an inverse

  std::unique_ptr<xla::PjRtClient> pjrt_client_;
  xla::PjRtDevice* device_;
  absl::BitGen bitgen_;
};

TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAndDequantizeToValidGraph) {
  constexpr absl::string_view kProgram = R"mlir(
func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
  %scale = "tf.Const"() { value = dense<0.347> : tensor<f32> } : () -> tensor<f32>

          TPUPartitionedOpConversionPass> {
  void runOnOperation() override;
};

template <typename T>
LogicalResult ReplacePartitionedOp(IntegerAttr num_cores_per_replica, T op) {
  constexpr bool is_input =
      std::is_same_v<std::decay_t<T>, TF::TPUPartitionedInputOp>;
  static_assert(
      is_input || std::is_same_v<std::decay_t<T>, TF::TPUPartitionedOutputOp>,

namespace mlir {
namespace TFTPU {
namespace {
constexpr char kStepMarkerLocationAttr[] = "step_marker_location";
constexpr char kUseXlaSpmdAttr[] = "use_spmd_for_xla_partitioning";

constexpr char kBadStringArrayElementMsg[] =
    "bad '{0}' attribute at index {1}, not a string";
constexpr char kBadArrayElementMsg[] =
    "bad '{0}' attribute at index {1} with value '{2}': failed to parse to {3}";
constexpr char kBadArrayAttrLengthMsg[] =