void IdentifyDilatedConvPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();

  patterns.add<ConvertTFDilatedConvOp<TF::Conv2DOp>,
               ConvertTFDilatedConvOp<TF::DepthwiseConv2dNativeOp>>(
      &getContext());
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}
}  // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateIdentifyDilatedConvPass() {

// TODO: b/289744814 - Refactor to have a single source of truth of TF Quant
// specs.
absl::flat_hash_set<int> GetQuantizableOperands(Operation* op) {
  absl::flat_hash_set<int> quantizable_operands;
  if (isa<TF::DepthwiseConv2dNativeOp, TF::Conv2DOp, TF::Conv3DOp, TF::MatMulOp,
          TF::BatchMatMulOp>(op)) {
    quantizable_operands.insert(1);
  } else if (isa<TF::GatherOp>(op)) {
    quantizable_operands.insert(0);

  // Supported quantizable ops.
  return isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::MatMulOp, TF::Conv2DOp,
             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp,
             TF::ResourceGatherOp, TF::DepthwiseConv2dNativeOp, TF::Conv3DOp,
             TF::BatchMatMulV2Op, TF::EinsumOp>(op);
}

bool IsOpWithInt8TypeOperand(Operation* op) {
  return (isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::XlaGatherOp, TF::GatherOp,