bcast_dims);
    } else {
      return rewriter.notifyMatchFailure(bcast_op, "Unsupported element type.");
    }
    Value new_const_op =
        rewriter.create<mhlo::ConstantOp>(bcast_op.getLoc(), result);
    rewriter.replaceOp(bcast_op, {new_const_op});
    return success();
  }
};

using FoldBroadcastInDimBeforeMulOp =

  Value add_op_result = op->getResult(0);

  // Broadcast bias value if unmatched with output shape.
  auto bcast_op = TryCast<BroadcastInDimOp>(bias_op.getDefiningOp(),
                                            /*name=*/"broadcast_in_dim_op");

  if (failed(bcast_op)) {
    bcast_op = TryCast<DynamicBroadcastInDimOp>(
        bias_op.getDefiningOp(),
        /*name=*/"dynamic_broadcast_in_dim_op");
  }

    auto round_val = rewriter.create<TF::RoundOp>(
        sc_op.getLoc(), float_avg_pool_op.getOutput());
    auto icast_op = rewriter.create<TF::CastOp>(
        sc_op.getLoc(), q_result_type.clone(qtype.getStorageType()), round_val);
    avg_pool_op.getResult().replaceAllUsesWith(icast_op.getResult());
    return success();
  }
};

// Applies quantization on the model in TF dialect.
class QuantizePass

    TypeID::get<TF::AnyOp>(),
    TypeID::get<TF::AvgPoolOp>(),
    TypeID::get<TF::AvgPool3DGradOp>(),
    TypeID::get<TF::AvgPoolGradOp>(),
    TypeID::get<TF::BatchToSpaceNDOp>(),
    TypeID::get<TF::BitcastOp>(),
    TypeID::get<TF::BroadcastToOp>(),
    TypeID::get<TF::CollectivePermuteOp>(),
    TypeID::get<TF::ComplexOp>(),
    TypeID::get<TF::ConcatV2Op>(),
    TypeID::get<TF::ConjOp>(),