if (!entry_function_attr) return std::nullopt;
    return entry_function_attr.getValue();
  } else {
    TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
    const auto f_attr = call_op.getFAttr().dyn_cast<FlatSymbolRefAttr>();
    if (!f_attr) return std::nullopt;
    return f_attr.getValue();
  }
}

class InsertCustomAggregationOpsPass
    : public PassWrapper<InsertCustomAggregationOpsPass,

 private:
  LogicalResult matchAndRewrite(TF::TPUPartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
    auto module_op = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

    auto f_name = f_attr.getValue();
    func::FuncOp float_func =

  static bool AllowDynamicRangeQuantizedOperand(
      Operation* quantized_op, const CustomMap& custom_op_map) {
    auto call_op = cast<TF::PartitionedCallOp>(quantized_op);
    StringRef function_name =
        call_op.getFAttr().cast<FlatSymbolRefAttr>().getValue();
    // The below can be generalized as there are more read-only ops added such
    // as slice.
    const bool is_gather = function_name.contains("gather");

  }

  LogicalResult matchAndRewrite(TF::PartitionedCallOp op,
                                PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
    StringRef function_name = f_attr.getValue();
    if (!function_name.starts_with(kDequantizeFunctionName)) return failure();

    llvm::SmallVector<Operation*> users(op->getUsers().begin(),

  auto spec = std::make_unique<OpQuantSpec>();
  if (auto call_op = dyn_cast<TF::PartitionedCallOp>(op)) {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }
    if (function_name.contains("depthwise_conv2d")) {
      spec->coeff_op_quant_dim[1] = 3;

    func::FuncOp ResolveFunc(::mlir::SymbolTableCollection* table) {
      if (table)
        return table->lookupNearestSymbolFrom<func::FuncOp>(*this, getFAttr());
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, getFAttr());
    }
    // TODO(b/204997177): Deprecate and remove.
    func::FuncOp func() {  return ResolveFunc(nullptr); }
  }];
}

            << min_num_elements_for_weights_ << " elements).";
        call_op->removeAttr(kQuantTraitAttrName);
      }
    }

    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if ((quantization_method_ == tensorflow::quantization::QuantizationMethod::
                                     METHOD_DYNAMIC_RANGE_INT8) &&
        (function_name.contains("batch_matmul") ||

  LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
    // removeAttr will return nullptr if no attribute was removed.
    if (!call_op->removeAttr(kQuantTraitAttrName) || !f_attr) {
      return failure();
    }

  }

  LogicalResult matchAndRewrite(TF::PartitionedCallOp op,
                                PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
    // Non-quantizable op
    if (!op->hasAttr(kQuantTraitAttrName)) return failure();
    StringRef function_name = f_attr.getValue();
    // TODO(b/228928859): Improve the getter function to match attributes rather

}

// Returns the function attribute for the given call op which is lifted for
// quantization.
inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
}

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}