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);

// Checks if an op is quantizable in StableHLO quantizer. Argument op is not
// necessarily a StableHLO op.
bool IsOpQuantizableStableHlo(Operation* op);

}  // namespace mlir::quant::stablehlo

    for (const auto [operand_idx, per_channel_dim] : spec->coeff_op_quant_dim) {
      spec->quantizable_operands.insert(operand_idx);
    }
  }
  return spec;
}

std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op) {
  auto scale_spec = std::make_unique<OpQuantScaleSpec>();
  if (llvm::isa<mlir::stablehlo::BroadcastInDimOp,
                mlir::stablehlo::ConcatenateOp,

    // For ops that require same operand and result types, use explicit
    // requantize op rather than using `entry_func_op`'s result as op result.
    auto spec = GetStableHloQuantConstraints(singular_op);
    const bool has_same_operand_and_result_type =
        spec->has_same_operand_and_result_type_requirement;
    if (has_same_operand_and_result_type) {

  ModuleOp module_op = getOperation();
  MLIRContext* ctx = module_op.getContext();

  auto func_op_quant_spec = GetStableHloOpQuantSpec;
  auto func_op_quant_scale_spec = GetStableHloQuantConstraints;

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // The function might contain more stats ops than required, and it will

      if (candidate_op->hasTrait<OpTrait::IsTerminator>()) {
        return failure();
      }

      if (!IsOpQuantizableStableHlo(candidate_op)) {
        return failure();
      }

      if (GetStableHloQuantConstraints(candidate_op)
              ->has_same_scale_requirement &&
          !IsConnectedWithQuantizedCompsiteFunction(candidate_op)) {
        return failure();
      }

  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  Operation* gather_op = FindOperationOfType<GatherOp>(main_fn);
  const auto spec = GetStableHloQuantConstraints(gather_op);

  EXPECT_THAT(spec, NotNull());
  EXPECT_THAT(spec->has_same_operand_and_result_type_requirement, IsTrue());
}

}  // namespace