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(

constexpr float kMinSparsityLevel = 0.3;
// Heuristic to check if a block configuration is correct for float constants.
constexpr float kBlockOverRandomSparsityRatio = 0.9;
// After quantization, some non-zero values are set to 0.
// Lower the ratio for identifying block configuration for quantized constants.
constexpr float kBlockOverRandomSparsityRatioQuant = 0.8;

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

namespace mlir {
namespace quant {
namespace {
constexpr StringRef kDequantizeFunctionName = "composite_dequantize";
constexpr StringRef kUniformQuantizationFunctionName = "uniform";

// Pre-actions before adding quantization logics. It creates a function with the

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kMiniBatchSplitsAttr[] = "mini_batch_splits";
constexpr char kMiniBatchCsrAttr[] = "mini_batch_in_csr";

#define GEN_PASS_DEF_EMBEDDINGPROGRAMKEYPASS

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

namespace mlir {
namespace TF {
namespace {

constexpr char kResourceFunctionMsg[] =
    "expects function level resource argument";
constexpr char kInvalidResourceMsg[] =
    "expects resource to be a VarHandleOp or function argument";
constexpr char kResourceNameArgAttr[] = "tf.resource_name";

// Checks if a function has only one block.

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

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kReplicaIdAttr[] = "_xla_replica_id";
constexpr char kDeviceOrdinalAttr[] = "device_ordinal";

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

using mlir::ValueRange;
using mlir::WalkResult;

constexpr llvm::StringRef kDeviceAttr = "device";
constexpr llvm::StringRef kNameAttr = "name";
constexpr llvm::StringRef kNumCoresPerReplicaAttr = "num_cores_per_replica";
constexpr llvm::StringRef kNumReplicasAttr = "num_replicas";
constexpr llvm::StringRef kMirroredVariableIndicesAttr =
    "_mirrored_variable_indices";

constexpr llvm::StringRef kBadReplicateInfoAttrMsg =

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

namespace mlir {
namespace TFTPU {

namespace {

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

// Checks if a function only contains a tf_executor.graph.
bool IsSupportedGraph(func::FuncOp func) {
  if (!llvm::hasSingleElement(func)) return false;

      DenseFPElementsAttr value_attr =
          mlir::cast<DenseFPElementsAttr>(op.getValue());
      new_values.reserve(value_attr.getNumElements());

      constexpr float kMaxFloat16Value = 65504.f;
      constexpr float kMinFloat16Value = -65504.f;

      for (auto value : value_attr.template getValues<float>()) {
        new_values.push_back(Eigen::half(

          /*dynamic_key=*/dynamic_key, op.getRecvKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
    }

    constexpr int64_t kDefaultCostEstimate = 1000000;
    rewriter.replaceOpWithNewOp<TF::XlaHostComputeOp>(
        op, op.getResultTypes(), op.getInputs(),
        /*ancestors=*/rewriter.getArrayAttr({}),