PatternRewriter &rewriter) const override {
    // Return early if the given operator is the custom aggregator op.
    if (dyn_cast_or_null<TF::CustomAggregatorOp>(op)) return failure();

    // The CustomAggregatorOp is only added after quantizable values.
    SmallVector<Value> quantizable_values;
    SmallVector<std::string> aggregator_ids;
    if (IsCallToQuantizableLiftedFunction(op)) {

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

  This function assumes _graph mode_ (used when legacy TF1 is used or when eager
  mode is explicitly disabled) when running the graph. This step is used in
  order to collect the statistics in CustomAggregatorOp for quantization using
  the representative dataset for the actual data provided for inference.

  Args:
    model_dir: Path to SavedModel directory.