%2 = "tf.If"(%0, %arg0, %arg1) {then_branch = @add, else_branch = @sub, output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = true} : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
  // CHECK: %1 = "tf.StatefulPartitionedCall"(%0, %arg1)
  // CHECK-SAME: f = @add
  // CHECK-SAME: _underscore_attr = "something"
  // CHECK-SAME: device = "noodle"

// cluster) ops to a separate parallel_execute region to run on CPU.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractOutsideCompilationPass();

// Create a pass that encapsulates StatefulPartitionedCallOp within a cluster.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateXlaClusterFormationPass();

// Creates a pass that marks unsupported ops in device cluster for outside
// compilation.

  llvm::SmallVector<SymbolUserOpInterface> noinline_pcall_ops,
      outermost_pcall_ops;
  if (mlir::failed(
          mlir::GetOpsOfTypeUntilMiss<mlir::TF::StatefulPartitionedCallOp,
                                      mlir::TF::PartitionedCallOp>(
              func, symtab, /*predicate*/ has_no_compile_device_type,
              /*hits*/ noinline_pcall_ops,

      SymbolTable::getSymbolUses(func, &module.getBodyRegion());
  if (!function_uses.has_value()) return false;
  for (auto& function_use : function_uses.value()) {
    if (!llvm::isa<TF::PartitionedCallOp, TF::StatefulPartitionedCallOp>(
            function_use.getUser())) {
      return false;
    }
  }
  return true;
}

// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.

    runOnOperation() {
  ModuleOp module_op = getOperation();

  func::FuncOp main_func = FindMainFuncOp(module_op);
  if (!main_func) return;

  // In case the model has tf.StatefulPartitionedCallOp or tf.PartitionedCallOp,
  // we recursively find called functions and process StableHLO ops in them.
  SmallVector<func::FuncOp> func_ops;
  func_ops.push_back(main_func);
  int stablehlo_func_id = -1;

#include "tensorflow/core/public/version.h"

namespace tensorflow {

TEST(RearrangeFunctionArgumentForFunctionTest, Basic) {
  FunctionDefLibrary fdl;
  {
    // Function for StatefulPartitionedCall's "f", If's
    // "then_branch"/"else_branch".
    // "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
    // "ret0" = "arg1"
    // "ret1" = "arg0"
    tensorflow::Scope s = tensorflow::Scope::NewRootScope();

// subcomputation in the TF/XLA bridge.
bool SupportsCommunicationComputation(Operation* op) {
  return isa<TF::IfRegionOp, TF::WhileRegionOp, TF::CaseRegionOp,
             TF::XlaCallModuleOp, TF::StatefulPartitionedCallOp,
             TF::PartitionedCallOp, TF::LegacyCallOp>(op);
}

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

      has_lower_as_multi_device_function_attr = lower.getValue();
    tensorflow::FunctionCallInlinePolicy policy =
        tensorflow::GetFunctionCallInlinePolicy(
            isa<PartitionedCallOp, StatefulPartitionedCallOp>(call),
            has_lower_as_multi_device_function_attr);

    if (policy == tensorflow::FunctionCallInlinePolicy::kMultiDevicePlacer)
      return MultiDeviceProcessInlinedCallBlocks(call, inlinedBlocks);

bool IsAlwaysAllowlistedOp(Operation *op) {
  return llvm::isa<
      // clang-format off
      // go/keep-sorted start
      TF::ConstOp,
      TF::IdentityOp,
      TF::PartitionedCallOp,
      TF::StatefulPartitionedCallOp
      // go/keep-sorted end
      // clang-format on
      >(op);
}

// LINT.IfChange
// The list of quantizable ops in the Legacy Integer mode.
ABSL_ATTRIBUTE_NOINLINE const std::set<std::string>

      callee.getContext(), callee.getFunctionType().getInputs(),
      llvm::to_vector<4>(new_return.getOperandTypes())));
  return success();
}

// Updates a PartitionedCallOp/StatefulPartitionedCallOp according to the
// resource-lifted new callee function in lifting_info.
template <typename CallOpType>
void UpdatePartitionedCallOpWithNewCallee(