diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
index 93e5cc22c30..54600faca4b 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
@@ -73,9 +73,8 @@ LogicalResult CollectAndGroupOutsideClusterOps(Block* block,
 }
 
 // Moves `cluster_ops` to associated `launch_op` body.
-void MoveOutsideClusterOpsToLaunchOp(
-    tf_device::LaunchOp launch_op,
-    const llvm::SmallVector<Operation*, 8>& cluster_ops) {
+void MoveOutsideClusterOpsToLaunchOp(tf_device::LaunchOp launch_op,
+                                     llvm::ArrayRef<Operation*> cluster_ops) {
   MLIRContext* context = launch_op.getContext();
   Operation* terminator = launch_op.GetBody().getTerminator();
 
@@ -123,7 +122,7 @@ void PropagateParallelExecuteReturnToReplicate(
 
 // Extracts all externally provided operands of `cluster_ops`.
 llvm::SmallSetVector<Value, 4> GetExternalOperands(
-    const llvm::SmallVector<Operation*, 8>& cluster_ops) {
+    llvm::ArrayRef<Operation*> cluster_ops) {
   llvm::SmallSetVector<Value, 4> external_values;
 
   for (Operation* op : cluster_ops) {
@@ -143,7 +142,7 @@ llvm::SmallSetVector<Value, 4> GetExternalOperands(
 
 // Extracts all externally used outputs of `cluster_ops`.
 llvm::SmallVector<Value, 4> GetExternalOutputs(
-    const llvm::SmallVector<Operation*, 8>& cluster_ops) {
+    llvm::ArrayRef<Operation*> cluster_ops) {
   llvm::SmallSetVector<Value, 4> external_outputs;
 
   for (Operation* op : cluster_ops) {
@@ -166,7 +165,7 @@ llvm::SmallVector<Value, 4> GetExternalOutputs(
 // as an operand.  If there are no external_inputs, set insertion point to first
 // cluster_op.
 void SetHostComputeInsertion(
-    OpBuilder* builder, const llvm::SmallVector<Operation*, 8>& cluster_ops,
+    OpBuilder* builder, llvm::ArrayRef<Operation*> cluster_ops,
     const llvm::SmallSetVector<Value, 4>& external_inputs) {
   if (external_inputs.empty()) builder->setInsertionPoint(cluster_ops.front());
   for (const auto& cluster_op : cluster_ops) {
@@ -183,9 +182,9 @@ void SetHostComputeInsertion(
 // using `communication_key`.
 TF::_HostComputeMlirOp CreateHostCompute(
     OpBuilder* builder, tf_device::ClusterOp tpu_cluster,
-    const llvm::SmallVector<Operation*, 8>& cluster_ops,
+    llvm::ArrayRef<Operation*> cluster_ops,
     const llvm::SmallSetVector<Value, 4>& inputs, llvm::ArrayRef<Value> outputs,
-    const std::string& communication_key) {
+    llvm::StringRef communication_key) {
   llvm::SmallVector<Type, 4> device_output_types;
   for (const auto& output : outputs)
     device_output_types.push_back(output.getType());
@@ -201,10 +200,9 @@ TF::_HostComputeMlirOp CreateHostCompute(
 
 void MoveOutsideCompiledOps(
     tf_device::ClusterOp tpu_cluster, llvm::StringRef outside_cluster_name,
-    tf_device::LaunchOp host_launch_op,
-    const llvm::SmallVector<Operation*, 8>& cluster_ops,
+    tf_device::LaunchOp host_launch_op, llvm::ArrayRef<Operation*> cluster_ops,
     const llvm::SmallSetVector<Value, 4>& external_inputs,
-    const llvm::SmallVector<Value, 4>& external_outputs) {
+    llvm::ArrayRef<Value> external_outputs) {
   if (external_inputs.empty() && external_outputs.empty()) {
     MoveOutsideClusterOpsToLaunchOp(host_launch_op, cluster_ops);
     return;