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

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

TEST_F(GetStableHloOpQuantSpecTest,

        spec.coeff_op_quant_dim[operand_idx] =
            quantized_type.dimension_specs().dimension();
      }
    }
  }
}

}  // namespace

std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op) {
  auto spec = std::make_unique<OpQuantSpec>();
  if (auto call_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op)) {
    auto entry_function =

    return success();
  }
};

void PrepareQuantizePass::runOnOperation() {
  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