using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;
using mlir::WalkResult;

constexpr char kReplicateSharding[] = "";
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";

#define GEN_PASS_DEF_TPUSHARDINGIDENTIFICATIONPASS

// in the input module op.
class PostCalibrationComponent : public Component {
 public:
  // Name of the post-training quantization post-calibration step. Used for
  // debugging purposes.
  static constexpr absl::string_view kName = "quant_ptq_post_calibration";

  explicit PostCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,

// Device coordinates are defined as (x, y, z, core), thus resulting in a rank 4
// topology.
constexpr int kTPUTopologyRank = 4;

constexpr char kDeviceTPUSystem[] = "TPU_SYSTEM";
constexpr char kDeviceTPU[] = "TPU";
constexpr char kTPUReplicatedCore[] = "TPU_REPLICATED_CORE";
constexpr char kTPUReplicatedHost[] = "TPU_REPLICATED_HOST";
constexpr char kBadIntArrayElementMsg[] =
    "bad '{0}' attribute at index {1}, not an int";

using ::tensorflow::Session;
using ::tensorflow::Status;
using ::tensorflow::StatusOr;
using ::tensorflow::Tensor;
using ::tensorflow::Var;

namespace {

constexpr char kResourceNameArgAttr[] = "tf.resource_name";
constexpr char kSavedModelArgAttr[] = "tf_saved_model.bound_input";

LogicalResult LiftVariablesFromSession(
    ModuleOp module, Session* session,
    const SmallSet<StringRef, 4>& resource_names) {

#include "tensorflow/core/public/session.h"

namespace tensorflow {
namespace quantization {

// Default MLIR dump file prefix for TensorFlow quantization passes.
inline constexpr absl::string_view kDefaultTfQuantMlirDumpFilePrefix =
    "tf_quant";

// Preprocesses the `module_op` for quantization. The preprocess steps include
// freezing the variables in the graph into constants. `is_inliner_run`

using ::tensorflow::quantization::OpSet;

constexpr absl::string_view kQuantizeCompositeFunctionsStepName =
    "_quantize_composite_functions";
constexpr StringRef kQuantizeFuncName = "quantize_i8";
constexpr StringRef kDequantizeFuncName = "dequantize_i8";
constexpr StringRef kAttrMapAttribute = "attr_map";
constexpr StringRef kQuantizedOpsAttribute = "tf_quant.quantized_ops";
constexpr StringRef kCompositeFuncPrefix = "composite_";

inline constexpr char kDebugModeOpQuantAttrName[] = "debug_quant";

// Used to annotate custom ops if they are quantizable.
inline constexpr char kQuantTraitAttrName[] = "_tfl_quant_trait";
enum QuantizationTrait { FullyQuantizable = 0, NotQuantizable = 1 };
inline constexpr absl::string_view QuantTraitValues[] = {"fully_quantizable",
                                                         "not_quantizable"};

inline constexpr double kNearZeroTolerance = 1.0e-6;

// 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: "

#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STRING_UTILS_H_

#include <stdint.h>

#include <cstddef>
#include <limits>
#include <vector>

#include "absl/status/status.h"

namespace mlir::TFL {

constexpr uint64_t kDefaultMaxLength = std::numeric_limits<int>::max();

class MiniDynamicBuffer {
 public:
  explicit MiniDynamicBuffer(size_t max_length = kDefaultMaxLength)
      : offset_({0}), max_length_(max_length) {}