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) {

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

}
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

  while (block_arg || cast_op_input) {
    if (block_arg) {
      // Change on device function type/shape.
      HandleFuncOp(block_arg.getOwner()->getParentOp());
      block_arg = nullptr;
      cast_op_input = nullptr;
    } else {
      auto cast_input = cast_op_input.getX();
      // Update input type.
      auto transform_result_type =

      Operation* def = input.getDefiningOp();
      DCHECK_EQ(def->getParentOp(), graph);
      try_update_current_candidate(def);
    }

    // Check island data operands.
    island.walk([&](Operation* op) {
      for (Value input : op->getOperands()) {
        Operation* def = input.getDefiningOp();
        if (!def || def->getParentOp() != graph) continue;

        try_update_current_candidate(def);
      }

// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {
  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}