func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {

    : (tensor<!tf_type.resource>, tensor<i1>) -> tensor<!tf_type.resource>
  // CHECK: "tf.PartitionedCall"
  // CHECK-SAME: f = @callee_stack_decomposed
  %call2 = "tf.PartitionedCall"(%stack, %arg0) {f = @callee, config = "", config_proto = "", executor_type = ""}
    : (tensor<!tf_type.resource>, tensor<i1>) -> tensor<!tf_type.resource>
  // CHECK: "tf.Slice"

    %cst_7 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %0 = "tf.PartitionedCall"(%arg0, %cst_1, %cst_2) {config = "", config_proto = "", executor_type = "", f = @quantize_i8} : (tensor<1x3x4x3xf32>, tensor<f32>, tensor<i32>) -> tensor<1x3x4x3xi8>

// PartitionedCall are placed on the host, a producer that produces them on the
// device will incur a D2H copy, even if the PartitionedCall is not executed
// (i.e. even if we choose to execute the XLA compiled computation via _XlaRun).
// To prevent this, we add control dependencies to make the int32 input edges
// into the PartitionedCall dead.  With this change the D2H copy only happens if

    func.return %2 : tensor<*xi32>
  }
}

// -----

// CHECK: module attributes {tfl._legalize_tfl_variables = true}
module {
  func.func @main() -> tensor<*xi32> {
    %0 = "tf.PartitionedCall"() {f = @f, config = "", config_proto = "", executor_type = ""}
      : () -> tensor<*xi32>
    func.return %0 : tensor<*xi32>
  }
  func.func @f() -> tensor<*xi32> {

    : (tensor<!tf_type.resource>) -> tensor<!tf_type.resource>
  // CHECK: "tf.PartitionedCall"(%[[VAR]], %[[GVAR1]], %[[GVAR2]])
  // CHECK-SAME: f = @callee_tensorarray_decomposed
  %call2 = "tf.PartitionedCall"(%call) {f = @callee, config = "", config_proto = "", executor_type = ""}
    : (tensor<!tf_type.resource>) -> tensor<!tf_type.resource>

  Option<bool> skip_partitioned_calls_{
      *this, "skip-partitioned-calls",
      ::llvm::cl::desc(
          "Skip tf.StatefulPartitionedCall and tf.PartitionedCall")};
};

void TFToMhloPass::runOnOperation() {
  auto func = getOperation();
  MLIRContext *context = func->getContext();

  RewritePatternSet patterns(context);

// CHECK-NEXT: [[q_bias:%.+]] = "quantfork.qcast"([[bias]]) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i32:f32, 0.044022349891595126>>

    up creating nested XLA launch ops.

    For example, the `tf.PartitionedCall` operation in the following code

    ```mlir
    func.func @convert_partitioned_call_with_resources(%arg0: tensor<!tf_type.resource>, %arg1: tensor<i32>) -> tensor<i32> {

  *library.add_function() = identity_func;

  Output in = ops::Placeholder(root, DT_FLOAT);
  NameAttrList b_name_attr;
  b_name_attr.set_name("IdentityFunc");
  ops::PartitionedCall call(root.WithOpName("call"), {in}, {DT_FLOAT},
                            b_name_attr);

  GraphDef graph_def;
  TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
  TF_ASSERT_OK(root.ToGraphDef(&graph_def));