graph = ops.get_default_graph()
    for op in graph.get_operations():
      if op.type == '_Retval':
        expected_node_name = op.name
        if op.get_attr('tf_saved_model.index_path') is not None:
          index_path_name = op.get_attr('tf_saved_model.index_path')[0]
          index_path_name = index_path_name.decode('utf-8').split(':')[0]
          try:

  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;

  auto func = builder.create<func::FuncOp>(module->getLoc(),
                                           kSessionInitFuncName, func_type);
  func->setAttr(kTfSavedModelExportedNamesAttr,
                builder.getStrArrayAttr({kSessionInitFuncName}));
  func->setAttr(kTfSavedModelInitializerTypeAttr,
                builder.getStringAttr(kTfSavedModelInitializerRestoreType));
  func.setVisibility(mlir::func::FuncOp::Visibility::Public);

  void runOnOperation() override {
    auto module = getOperation();
    // If TFLite variable legalization is not allowed, then we skip this pass.
    if (auto legalize_tfl_variables_attr =
            module->getAttr(kLegalizeTflVariables)) {
      if (!mlir::cast<BoolAttr>(legalize_tfl_variables_attr).getValue()) return;
    }

    RewritePatternSet patterns(&getContext());
    populateWithGenerated(patterns);

  builder.setInsertionPoint(execute_launch);
  auto copy_with_layout = BuildCopyWithLayout(execute_launch, compile_launch,
                                              get_layout, input, &builder);
  copy_with_layout->setAttr(kDeviceAttr, execute_launch.getDeviceAttr());
  execute.setOperand(execute_arg_index, copy_with_layout);
}

// Performs transformation for replicated inputs. Returns true if this is a

  new_call_op->setAttr(kEntryFuncAttrName,
                       rewriter.getStringAttr(new_ref_func_name.getValue()));
  new_call_op->setAttrs(call_op->getAttrs());
  new_call_op->removeAttr(rewriter.getStringAttr(kQuantTraitAttrName));

  FlatSymbolRefAttr new_func_name_attr =
      FlatSymbolRefAttr::get(rewriter.getContext(), new_ref_func_name);
  new_call_op->setAttr(kEntryFuncAttrName, new_func_name_attr);

  }

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

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

 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 =

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