TF_RETURN_IF_ERROR(c->GetAttr("num_bins", &num_bins_attr));
      c->set_output(3, c->MakeShape({num_bins_attr->i()}));

      return absl::OkStatus();
    });

class CustomAggregatorOp : public OpKernel {
 public:
  explicit CustomAggregatorOp(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("id", &id_));

    int calibration_method_value;
    int num_bins;

    : public OpRewritePattern<TF::CustomAggregatorOp> {
 public:
  // Does not take ownership of context, which must refer to a valid value that
  // outlives this object.
  explicit ConvertCustomAggregationOpToQuantStats(MLIRContext *context)
      : OpRewritePattern<TF::CustomAggregatorOp>(context) {}

  LogicalResult matchAndRewrite(TF::CustomAggregatorOp op,

                            StartsWith("composite_dot_general_fn")));

  // Tests that there are CustomAggregatorOps inserted.
  SmallVector<TF::CustomAggregatorOp> custom_aggregator_ops;
  for (auto custom_aggregator_op :
       func_ops[0].getOps<TF::CustomAggregatorOp>()) {
    custom_aggregator_ops.push_back(custom_aggregator_op);
  }
  EXPECT_THAT(custom_aggregator_ops, SizeIs(2));
}

}  // namespace

namespace mlir::quant::stablehlo {

// Performs pre-calibration graph transformation as part of post-training
// static-range quantization.

// The resulting `ModuleOp` contains `TF::CustomAggregatorOp`s for collecting
// quantization statistics, along with `TF::XlaCallModuleOp`s that correspond to
// lifted quantizable functions.
class PreCalibrationComponent : public Component {
 public:

//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.
class PostCalibrationComponent : public Component {
 public:
  // Name of the post-training quantization post-calibration step. Used for
  // debugging purposes.

                      ReadStatistics(calibration_data_dir));

  absl::Status status = absl::OkStatus();
  module_op.walk([&py_function_library, &calibration_options, &status,
                  &statistics_map](mlir::TF::CustomAggregatorOp aggregator_op) {
    mlir::StringRef id = aggregator_op.getId();
    auto iter = statistics_map.find(id);
    if (iter == statistics_map.end()) {
      status = absl::InternalError(

//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.
//
// TODO: b/320607042 - Add tests for this component on the python layer.
class CalibrationComponent : public Component {
 public:

    const std::unordered_set<std::string>& aggregator_ops_to_ignore,
    SmallVector<Value>& statistics_outputs, SmallVector<StringRef>& ids,
    SmallVector<int32_t>& calibration_methods) {
  for (auto op : region.getOps<TF::CustomAggregatorOp>()) {
    if (aggregator_ops_to_ignore.count(op.getId().str())) continue;

    ids.push_back(op.getId());
    calibration_methods.push_back(op.getCalibrationMethod());
    statistics_outputs.push_back(op.getMin());