// CHECK: [[results_0:%.*]]:2 = tfrt.call @callee([[chain]]
  // CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
  %1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_1:%.*]]:2 = tfrt.call @callee([[chain]]

    return llvm::isa_and_nonnull<mlir::TF::StatefulPartitionedCallOp>(op)
               ? mlir::WalkResult::interrupt()
               : mlir::WalkResult::advance();
  });
  if (result.wasInterrupted()) {
    // StatefulPartitionedCall ops are not supported by the tflite runtime.
    mlir::TFL::AttachErrorCode(
        module.emitError(
            "The Graph contains unsupported `StatefulPartionedCallOp`(s), will "

  // TPUCompilationResultsOp.
  if (op->hasAttr(kTpuCompilationStatus)) {
    op->setAttr(kTpuCompilationStatus, attr);
  }
}

// A helper class to inline TF::StatefulPartitionedCall ops
struct Inliner : public InlinerInterface {
  Inliner(OpBuilder& builder, SymbolTable& symbol_table)
      : InlinerInterface(builder.getContext()),
        builder(builder),

      result.replaceAllUsesWith(if_op.getOperand(then_aliased_input + 1));
    }
  }
  if_op.replaceAllUsesWith(new_if);
  if_op.erase();
  return success();
}

// Handles stack usage by a tf.StatefulPartitionedCall or a tf.PartitionedCall.
// It will first check if the callee was previously handled, and try to reuse
// that result if so. Otherwise, it will clone and convert the callee function,

  // CHECK: [[results_0:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
  %1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_1:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor

                                                arg_buffer_size_is_fixed(i)};
    }
  }
  UpdateFuncType(func);
}

// Holds information about a decomposed callee function for
// PartitionedCall/StatefulPartitionedCall.
struct PartitionedCallDecompositionInfo {
  bool signature_change;
  func::FuncOp decomposed_callee;
  llvm::SmallDenseMap<int64_t, int64_t> buffer_arg_to_size_arg;

    // (checked later).
    if (auto device_type_attr =
            op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr)) {
      // tf.StatefulPartitionedCall ops with and without
      // _tpu_replicate attributes may exist in the same graph. Ops without
      // the attribute but with _XlaMustCompile=true would have
      // _xla_compile_device_type="" after

    %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @partitioned_called_func} : (tensor<20xi32>) -> tensor<*xi32>
    // CHECK: tf.StatefulPartitionedCall
    // CHECK-SAME: (tensor<20xi32>) -> tensor<20xi32>
    %1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @stateful_partitioned_call_func} : (tensor<20xi32>) -> tensor<*xi32>
    // CHECK: tf.TPUPartitionedCall

func.func @func_with_device_training_loop(%arg0: tensor<*xi32>, %arg1: tensor<*xi1>) -> (tensor<*xi32>, tensor<*xi1>) {
  %1:2 = "tf.StatefulPartitionedCall"(%arg0){f= @func_body, config="", config_proto="", executor_type=""}
         : (tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi1>)

  // CHECK: call @callee() : () -> ()
  // CHECK: call @callee() : () -> ()
  "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @callee} : () -> ()
  "tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @callee} : () -> ()
  "tf.LegacyCall"() {config = "", config_proto = "", executor_type = "", f = @callee} : () -> ()
  func.return
}