From cb4c7d28f0f1bb2103b6de6c6000909a78417370 Mon Sep 17 00:00:00 2001
From: Xiao Yu <fishx@google.com>
Date: Fri, 29 Jan 2021 11:15:31 -0800
Subject: [PATCH] Fix a flaky mlir test.

PiperOrigin-RevId: 354570431
Change-Id: I0cde67a9487888e7e285462b17fe371f2231a736
---
 .../mlir/tensorflow/tests/cluster_tf_ops_pass.mlir   | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir b/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir
index fd3519a4c8c..2ae5ebc518b 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir
@@ -9,7 +9,7 @@
 // CHECK-SAME:   body = @while_body
 // CHECK-SAME:   cond = @while_cond
 // CHECK-SAME:   device = "/job:localhost/replica:0/task:0/device:CPU:0"
-// CHECK-NEXT:   %[[RESULT_1:.*]] = tf_device.remote_run "/job:worker/replica:0/task:1" @_job_worker_replica_0_task_1(%[[ARG_1]])
+// CHECK-NEXT:   %[[RESULT_1:.*]] = tf_device.remote_run "/job:worker/replica:0/task:1" @[[MAIN_PARTITION_0:.*]](%[[ARG_1]])
 // CHECK-NEXT:   return %[[RESULT_0]], %[[RESULT_1]]
 func @main(%arg0: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg1: tensor<i32> {tf.device = "/job:worker/replica:0/task:1/device:CPU:0"}) -> (tensor<i32>, tensor<i32>) {
   %1 = "tf.While"(%arg0) {cond = @while_cond, body = @while_body, is_stateless = false, shape_invariant, device="/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> (tensor<i32>)
@@ -40,9 +40,9 @@ func @while_cond(%arg0: tensor<i32> {tf.device = "/job:localhost/replica:0/task:
 // CHECK-NEXT:   %[[RESULT_2:.*]] = "tf.Const"() {value = dense<16> : tensor<i32>} : () -> tensor<i32>
 // CHECK-NEXT:   tf_device.send %[[RESULT_2]] "key-0" "/job:worker/replica:0/task:1/device:CPU:0"
 // CHECK-SAME:  device = "/job:localhost/replica:0/task:0/device:CPU:0"
-// CHECK-NEXT:   tf_device.remote_run "/job:worker/replica:0/task:1" @_job_worker_replica_0_task_1_0()
+// CHECK-NEXT:   tf_device.remote_run "/job:worker/replica:0/task:1" @[[BODY_PARTITION_0:.*]]() : () -> ()
 // CHECK-NEXT:   tf_device.send %[[RESULT_2]]
-// CHECK-NEXT:   tf_device.remote_run "/job:worker/replica:0/task:2" @_job_worker_replica_0_task_2()
+// CHECK-NEXT:   tf_device.remote_run "/job:worker/replica:0/task:2" @[[BODY_PARTITION_1:.*]]() : () -> ()
 // TODO(tf-runtime): Allow while body having remote inputs and outputs.
 func @while_body(%arg0: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<i32>) {
   %0 = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
@@ -59,20 +59,20 @@ func @while_body(%arg0: tensor<i32> {tf.device = "/job:localhost/replica:0/task:
   return %1 : tensor<i32>
 }
 // Subgraph of @main function that is placed on worker:1
-// CHECK: func @_job_worker_replica_0_task_1(%[[ARG_0:.*]]: tensor<i32> {tf.device = "/job:worker/replica:0/task:1/device:CPU:0"})
+// CHECK: func @[[MAIN_PARTITION_0]](%[[ARG_0:.*]]: tensor<i32> {tf.device = "/job:worker/replica:0/task:1/device:CPU:0"})
 // CHECK-SAME:  host = "/job:worker/replica:0/task:1"
 // CHECK-NEXT:   %[[RESULT_0:.*]] = "tf.AddV2"(%[[ARG_0]], %[[ARG_0]])
 // CHECK-SAME:  device = "/job:worker/replica:0/task:1/device:CPU:0"
 // CHECK-NEXT:   return %[[RESULT_0]]
 
 // Subgraph of @while_body function that is placed on worker:1
-// CHECK: func @_job_worker_replica_0_task_1_0() attributes {host = "/job:worker/replica:0/task:1"}
+// CHECK: func @[[BODY_PARTITION_0]]() attributes {host = "/job:worker/replica:0/task:1"}
 // CHECK-NEXT:   %[[RESULT_0:.*]] = tf_device.receive "key-0"
 // CHECK-NEXT:   %[[RESULT_1:.*]] = "tf.AddV2"(%[[RESULT_0]], %[[RESULT_0]])
 // CHECK-SAME:  device = "/job:worker/replica:0/task:1/device:CPU:0"
 
 // Subgraph of @while_body function that is placed on worker:2
-// CHECK: func @_job_worker_replica_0_task_2() attributes {host = "/job:worker/replica:0/task:2"}
+// CHECK: func @[[BODY_PARTITION_1]]() attributes {host = "/job:worker/replica:0/task:2"}
 // CHECK-NEXT:   %[[RESULT_0:.*]] = tf_device.receive "key-1"
 // CHECK-NEXT:   %[[RESULT_1:.*]] = "tf.AddV2"(%[[RESULT_0]], %[[RESULT_0]])
 // CHECK-SAME:  device = "/job:worker/replica:0/task:2/device:CPU:0"