if (llvm::isa<TF::_TPUCompileMlirOp>(op)) {
      op->setAttr("_is_compilation", UnitAttr::get(func->getContext()));
      op = op->getParentOp();
      while (op && op != func) {
        op->setAttr("_wraps_compilation", UnitAttr::get(func->getContext()));
        op = op->getParentOp();
      }
    }
  });
}

void UnmarkCompilationOps(Operation* func) {
  func->walk([&](Operation* op) {

  bool in_launch = isa<tf_device::LaunchOp>(preprocess_op->getParentOp());
  Operation* preprocess_or_launch =
      in_launch ? preprocess_op->getParentOp() : preprocess_op;

  Operation* tpu_compile_successor = nullptr;
  func_op->walk([&](TF::_TPUCompileMlirOp compile_op) {
    if (compile_op->getParentOp() == nullptr ||
        !isa<tf_device::LaunchOp>(compile_op->getParentOp()))
      return WalkResult::advance();

    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// Check if the TPUCopyWithDynamicShapeOp is valid.
// 1. The op should be wrapped inside a launch op.
// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {
  auto launch_op = llvm::dyn_cast<tf_device::LaunchOp>(op->getParentOp());
  if (!launch_op) {

  EXPECT_EQ(noinline_pcall_ops.size(), 2);
  auto func =
      llvm::dyn_cast<mlir::func::FuncOp>(noinline_pcall_ops[0]->getParentOp());
  ASSERT_TRUE(func);
  EXPECT_EQ(func.getSymName(), "while_body_func");
  func =
      llvm::dyn_cast<mlir::func::FuncOp>(noinline_pcall_ops[1]->getParentOp());
  ASSERT_TRUE(func);
  EXPECT_EQ(func.getSymName(), "outer_stateful_pcall_func");

  EXPECT_EQ(outermost_pcall_ops.size(), 1);

        // For a block argument, consider transforms only when it is a
        // replicated input (defining ops will be outside the replicate node).
        if (maybe_replicate != block_arg.getParentRegion()->getParentOp() ||
            !HandleReplicatedInputs(execute_arg_index, execute.getKey(),
                                    execute_launch, compile_launch, block_arg,

  void runOnOperation() override;
};

// Get the ancestor of `descendant` that is a direct child of `ancestor`.
Operation* GetAncestorBelow(Operation* descendant, Operation* ancestor) {
  Operation* parent = descendant->getParentOp();
  if (!parent) return nullptr;
  if (parent == ancestor) return descendant;
  return GetAncestorBelow(parent, ancestor);
}

// `is_cpu_read` is set to `true` iff `read` is on a resource with device type
// CPU.

      const bool parent_is_replicate =
          llvm::isa<tf_device::ReplicateOp>(execute_launch->getParentOp()) ||
          (llvm::isa<tf_device::ParallelExecuteOp>(
               execute_launch->getParentOp()) &&
           llvm::isa<tf_device::ReplicateOp>(
               execute_launch->getParentOp()->getParentOp()));

      // If this is inside a tf_device::ReplicateOp, the variables are

}
ResourceConstructingOps ResourceConstructingOps::EntryState(Value value) {
  if (auto barg = mlir::dyn_cast<BlockArgument>(value)) {
    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());
      auto global_tensor = tf_saved_model::LookupBoundInputOfType<
          tf_saved_model::GlobalTensorOp>(func, barg.getArgNumber(),

  // Use tf_executor.island op if present as non v1 control flow op results are
  // forwarded by a parent tf_executor.island op.
  if (llvm::isa<tf_executor::IslandOp>(op_to_update->getParentOp()))
    op_to_update = op_to_update->getParentOp();

  for (Value result : op_to_update->getResults()) {
    if (mlir::isa<tf_executor::TokenType>(result.getType())) continue;
    if (mlir::isa<tf_executor::ControlType>(result.getType())) break;

      getElementTypeOrSelf(value.getType()));
}

bool HasOutsideCompiledAncestor(Operation* op) {
  Operation* parent = op->getParentOp();
  while (parent) {
    if (parent->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr))
      return true;
    parent = parent->getParentOp();
  }
  return false;
}

// If any tf.variants are inputs/outputs to the another outside compiled