}
      return WalkResult::advance();
    });

    if (!walk_result.wasInterrupted())
      head_outside_compiled_ops.insert(&cluster_op);
  }

  return head_outside_compiled_ops.takeVector();
}

// Moves head outside compiled ops into its own `tf_device.LaunchOp`
// computation before the cluster.
void CreateHeadComputation(OpBuilder* builder,

      new_launch_op_results.insert(operand);
    }
  }

  for (Value result : tpu_copy_with_dynamic_shape_op->getResults()) {
    new_launch_op_results.remove(result);
  }

  return new_launch_op_results.takeVector();
}

// Create a new host launch op which contains all the old launch op body
// except the dynamic shape copy op.
tf_device::LaunchOp CreateNewHostLaunchOpWithNewResult(
    tf_device::LaunchOp* old_launch_op,

      }

      if (ops_to_reorder.insert(op)) {
        llvm::append_range(worklist, op->getResults());
      }
    }
  }

  llvm::SmallVector<mlir::Operation*, 0> sorted = ops_to_reorder.takeVector();
  llvm::sort(sorted, [](mlir::Operation* lhs, mlir::Operation* rhs) {
    return lhs->isBeforeInBlock(rhs);
  });

  for (mlir::Operation* const op : llvm::reverse(sorted)) {
    op->moveAfter(cluster);