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;

        "//tensorflow/core:test",
        "//tensorflow/core:test_main",
        "@com_google_googletest//:gtest_main",
    ],
)

tf_kernel_library(
    name = "custom_aggregator_op",
    srcs = ["custom_aggregator_op.cc"],
    compatible_with = get_compatible_with_portable(),
    visibility = [
        "//tensorflow:__pkg__",
        "//tensorflow/compiler/mlir/quantization/tensorflow/python:__pkg__",

  // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "0", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>

// CHECK: %[[CUSTOM_AGGREGATOR_1:.+]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE]]) {{.*}} : (tensor<1x4x4x8xf32>) -> (tensor<1x4x4x8xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)

// CHECK-LABEL: @serving_default
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)

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

}
// CHECK-LABEL: @serving_default
// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "0", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)

    : 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,

                      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: