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

namespace mlir::quant {

std::unique_ptr<OpQuantSpec> GetUniformOpQuantSpec(Operation* op) {
  auto spec = std::make_unique<OpQuantSpec>();
  if (isa<TF::UniformQuantizedConvolutionHybridOp>(op) ||
      isa<TF::UniformQuantizedConvolutionOp>(op)) {
    spec->coeff_op_quant_dim[1] = 3;

// is set for `spec`.
// TODO: b/323478683 - Duplicate tracking of config will be eliminated.
// `OpQuantSpec` will be deprecated and `Method` will be used instead.
void PopulateCoeffOpQuantDimIfPerChannelQuantized(
    TF::XlaCallModuleOp xla_call_module_op, OpQuantSpec& spec) {
  absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
  if (method.ok() && method->has_static_range_ptq()) {

  // name.
  std::vector<Record *> defs = records.getAllDerivedDefinitions("Op");
  llvm::sort(defs, LessRecord());

  OUT(0) << "static std::unique_ptr<quant::OpQuantSpec> "
            "GetOpQuantSpec(mlir::Operation *op, bool "
            "disable_per_channel_for_dense_layers = false) {\n";
  // TODO(b/176258587): Move to OpTrait if this should be generalized.

namespace mlir {
namespace quant {

// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetUniformOpQuantSpec(Operation* op);

}  // namespace quant
}  // namespace mlir

      return cur_spec;
  }
  return std::nullopt;
}

// TODO(b/228928859): Improve the getter function to match attributes rather
// than function name.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op) {
  auto spec = std::make_unique<OpQuantSpec>();
  if (auto call_op = dyn_cast<TF::PartitionedCallOp>(op)) {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();

    }
  }
)mlir";

// TOOD: b/323478683 - Directly use types rather than creating a `unique_ptr`.
std::unique_ptr<quant::OpQuantSpec> GetOpQuantSpec(
    const mlir::Operation* op,
    bool disable_per_channel_for_dense_layers = false) {
  auto spec = std::make_unique<quant::OpQuantSpec>();
  spec->coeff_op_quant_dim[1] = 3;
  spec->biases_params[2] = {{0, 1}, quant::GetUniformQuantizedTypeForBias};

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace mlir::quant::stablehlo {

// Returns StableHLO quantization specs for an op.
std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op);

// Returns quantization constraints (ex: fixed output, same scale) given
// a StableHLO op.
std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op);

GetWeightComponentSpec(
    const tensorflow::quantization::QuantizationOptions& quantization_options);

// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op);

// Returns quantization scale specs (fixed output, same scale) for a TF op.
std::unique_ptr<OpQuantScaleSpec> GetTfQuantScaleSpec(Operation* op);

}  // namespace quant

    // Check whether dynamic range quantization can be applied.
    for (auto& use : value.getUses()) {
      Operation* user = use.getOwner();
      int operand_num = use.getOperandNumber();
      std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(user);

      if (quant_specs_.inference_type == tensorflow::DT_QINT8 &&
          spec->quantizable_operands.contains(operand_num)) {
        quantizable_ops.insert({user, operand_num});

  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
  ASSERT_TRUE(succeeded(xla_call_module_op));

  const std::unique_ptr<OpQuantSpec> op_quant_spec =
      GetStableHloOpQuantSpec(*xla_call_module_op);
  ASSERT_THAT(op_quant_spec, NotNull());

  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim, IsEmpty());
}