class PropagateDequantizeOpIfAllowed
    : public OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit PropagateDequantizeOpIfAllowed(MLIRContext* context)
      : OpRewritePattern<TF::PartitionedCallOp>(context) {}

  // Create a new dequantize op that is propagated.
  void createNewDequantizeOp(PatternRewriter& rewriter,
                             TF::PartitionedCallOp original_dequantize_op,

  return quantization_func;
}

// Post-actions after adding quantization logics. Post-actions include
// 1) Adding the created function in the symbol table
// 2) Creating a PartitionedCallOp in the main graph that calls the created
//    function.
TF::PartitionedCallOp FinalizeFunctionRegister(
    PatternRewriter& rewriter, Value input, Value output,
    func::FuncOp& quantization_func, Operation* quantized_op,

}

Operation *DuplicateOp(TF::PartitionedCallOp call_op, PatternRewriter &rewriter,
                       const StringAttr &new_ref_func_name) {
  // Create PartitionedCallOp to the copied composite function. This
  // PartitionedCallOp does not have kQuantTraitAttrName, and therefore won't
  // get quantized.
  auto new_call_op = rewriter.create<TF::PartitionedCallOp>(

      xla_call_module_op != nullptr) {
    absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
    if (method.ok() && method->has_static_range_ptq()) return true;
  }

  TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
  return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
             llvm::StringRef(

class PreprocessConstantOp : public OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit PreprocessConstantOp(MLIRContext* context, OpSet op_set,
                                QuantMethod quantization_method,
                                bool enable_per_channel_quantization)
      : OpRewritePattern<TF::PartitionedCallOp>(context),
        op_set_(op_set),
        quantization_method_(quantization_method),

// into [H, W, In, Mul]
class RestoreWeightShapePattern
    : public OpRewritePattern<TF::PartitionedCallOp> {
  using OpRewritePattern<TF::PartitionedCallOp>::OpRewritePattern;

 private:
  LogicalResult addReshapeOpToDepthwiseWeight(TF::PartitionedCallOp op,
                                              PatternRewriter& rewriter) const {
    int weight_operand_idx = 1;

class CheckQuantizableOps
    : public mlir::OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit CheckQuantizableOps(MLIRContext* context,
                               const QuantizationOptions& quant_options)
      : OpRewritePattern<TF::PartitionedCallOp>(context),
        quant_options_(quant_options) {}

 private:
  LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,

      if (!preceding_op) continue;

      // Check whether the preceding op is a quantized composite function.
      if (llvm::isa<TF::PartitionedCallOp>(preceding_op)) {
        auto call_op = llvm::cast<TF::PartitionedCallOp>(preceding_op);
        if (!IsCompositeFunction(call_op)) continue;
        return true;
      }

      // Check if the preceding op is a quantized same-scale op.

namespace TFDevice {

namespace {

#define GEN_PASS_DEF_CONVERTLAUNCHFUNCTOTFCALLPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

// Rewrites tf_device::LaunchFuncOp into TF::PartitionedCallOp.
struct ConvertLaunchFuncToTFCallPass
    : public impl::ConvertLaunchFuncToTFCallPassBase<
          ConvertLaunchFuncToTFCallPass> {
  void runOnOperation() override;
};

  // In graph mode, we create a PartitionedCallOp instead:
  // https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L573

  // TODO(bmzhao): After discussing with Allen, we should execute this via a
  // PartitionedCallOp for compatibility with "tooling that assumes functions in
  // graphs are PartitionedCallOps".