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,

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

class CheckQuantizableOps
    : public mlir::OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit CheckQuantizableOps(MLIRContext* context,
                               const QuantMethod quantization_method,
                               const OpSet target_opset,
                               const int min_num_elements_for_weights)
      : OpRewritePattern<TF::PartitionedCallOp>(context),
        quantization_method_(quantization_method),

// skipped from outlining for both single-core and replicated case (i.e., the
// `_tpu_v1_compat_outlined` module must be empty and no `PartitionedCallOp` is
// created). Also check that `_skip_island_outlining` attribute is removed.

// CHECK-LABEL: @func5
// CHECK-NOT: _skip_island_outlining
// CHECK-NOT: PartitionedCallOp
// CHECK: _tpu_v1_compat_outlined {
// CHECK-NEXT: }
func.func @func5() attributes {_skip_island_outlining = true} {

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

std::optional<TF::PartitionedCallOp> ApplyUniformQuantization(
    PatternRewriter& rewriter, TF::ConstOp op,
    tensorflow::quantization::QuantizationComponentSpec& weight_spec);

}  // namespace quant
}  // namespace mlir

    %0 = "tf.some_op"(%arg0) : (tensor<i32>) -> tensor<i32>
    func.return %0 : tensor<i32>
  }
  func.func @while_cond_without_cluster_attr(%arg0: tensor<i32>) -> tensor<i1> {
    %0 = "tf.PartitionedCallOp"(%arg0) { f = @callee_func} : (tensor<i32>) -> tensor<i1>
    func.return %0 : tensor<i1>
  }
  func.func @callee_func(%arg0: tensor<i32>) -> tensor<i1> {
    %0 = "tf.some_op"(%arg0) : (tensor<i32>) -> tensor<i1>

  quant_spec.set_quantization_component(
      QuantizationComponentSpec::COMPONENT_WEIGHT);
  quant_spec.set_tensor_type(QuantizationComponentSpec::TENSORTYPE_INT_8);

  std::optional<TF::PartitionedCallOp> dequantize_op = ApplyUniformQuantization(
      pattern_rewriter, cast<TF::ConstOp>(value.getDefiningOp()), quant_spec);
  EXPECT_TRUE(dequantize_op.has_value());
  EXPECT_EQ(dequantize_op.value().func().getName().str(),

  Operation *nested_module = symbol_table.lookup(kNestedModule);
  if (!nested_module) return;

  InlinerInterface inliner(&getContext());
  auto walk_result = getOperation().walk([&](TF::PartitionedCallOp call_op) {
    if (!call_op.getF().getRootReference().getValue().starts_with(
            kNestedModule))
      return WalkResult::advance();
    // This is a call we need to inline!
    LLVM_DEBUG(llvm::dbgs()

==============================================================================*/

include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"

// Creates a function call with TF::PartitionedCallOp and a new function to
// wrap the section between arguments and results.
//
// The `returns` number indicates the number of results the function returns.
class LiftAsTFPartitionedCall<string func_name, int returns = 1> :