//tensorflow/dtensor/mlir/tests:spmd_expansion.mlir.test is flaky
@rishikasinha-tf 
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

git HEAD

### Custom code

No

### OS platform and distribution

Ubuntu 20.04

### Mobile device

n/a

### Python version

3.9.16

### Bazel version

6.1.0

### GCC/compiler version

10.2.1

### CUDA/cuDNN version

n/a

### GPU model and memory

n/a

### Current behavior?

//tensorflow/dtensor/mlir/tests:spmd_expansion.mlir.test unit test fails occasionally

x86 log
https://source.cloud.google.com/results/invocations/88910868-abeb-4aa0-8954-df2b79ef5a26/log

AARCH64 log
https://github.com/tensorflow/tensorflow/actions/runs/5420661339/jobs/9855097285

### Standalone code to reproduce the issue

```shell
bazel --bazelrc=/usertools/cpu.bazelrc test --config=rbe --config=pycpp --config=build_event_export
```


### Relevant log output

```shell
FAIL: //tensorflow/dtensor/mlir/tests:spmd_expansion.mlir.test (see /root/.cache/bazel/_bazel_root/fbac33eb30dbfb6b11b15a7ff5ac830d/execroot/org_tensorflow/bazel-out/k8-opt/testlogs/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test/test.log)
INFO: From Testing //tensorflow/dtensor/mlir/tests:spmd_expansion.mlir.test:
==================== Test output for //tensorflow/dtensor/mlir/tests:spmd_expansion.mlir.test:
-- Testing: 1 tests, 1 workers --
FAIL: MLIR tests :: spmd_expansion.mlir (1 of 1)
******************** TEST 'MLIR tests :: spmd_expansion.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';   /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/dtensor-opt -- /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics | /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/llvm-project/llvm/FileCheck /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir
--
Exit Code: 1

Command Output (stderr):
--
2023-07-03 09:20:51.133345: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
TensorFlow crashed, please file a bug on https://github.com/tensorflow/tensorflow/issues with the trace below.
Stack dump:
0.	Program arguments: /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/dtensor-opt /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/dtensor-opt /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics
1.	Program arguments: /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/dtensor-opt /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir -split-input-file -dtensor-annotate-global-shape -dtensor-spmd-expansion -verify-diagnostics
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  libtensorflow_framework.so.2 0x00007f49dfe75e6e llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 46
1  libtensorflow_framework.so.2 0x00007f49dfe73f07 llvm::sys::RunSignalHandlers() + 87
2  libtensorflow_framework.so.2 0x00007f49dfe76692
3  libpthread.so.0              0x00007f49dea05420
4  dtensor-opt                  0x0000562b92fdd978
5  dtensor-opt                  0x0000562b92fdd11a
6  dtensor-opt                  0x0000562b92fde2b9
7  dtensor-opt                  0x0000562b92f94499
8  dtensor-opt                  0x0000562b92fb7b91
9  dtensor-opt                  0x0000562b92fb85bf
10 dtensor-opt                  0x0000562b92eca421
11 dtensor-opt                  0x0000562b94e800d4
12 dtensor-opt                  0x0000562b94e82d17
13 dtensor-opt                  0x0000562b94e82ae5
14 dtensor-opt                  0x0000562b948df256
15 dtensor-opt                  0x0000562b948de53a
16 dtensor-opt                  0x0000562b9502d603
17 dtensor-opt                  0x0000562b9502d34e
18 dtensor-opt                  0x0000562b948da857
19 dtensor-opt                  0x0000562b948dad3d
20 dtensor-opt                  0x0000562b92d80fcc
21 libc.so.6                    0x00007f49de626083 __libc_start_main + 243
22 dtensor-opt                  0x0000562b92d80d39
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  libtensorflow_framework.so.2 0x00007f49dfe75e6e llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) + 46
1  libtensorflow_framework.so.2 0x00007f49dfe73f37 llvm::sys::RunSignalHandlers() + 135
2  libtensorflow_framework.so.2 0x00007f49dfe76692
3  libpthread.so.0              0x00007f49dea05420
4  dtensor-opt                  0x0000562b92fdd978
5  dtensor-opt                  0x0000562b92fdd11a
6  dtensor-opt                  0x0000562b92fde2b9
7  dtensor-opt                  0x0000562b92f94499
8  dtensor-opt                  0x0000562b92fb7b91
9  dtensor-opt                  0x0000562b92fb85bf
10 dtensor-opt                  0x0000562b92eca421
11 dtensor-opt                  0x0000562b94e800d4
12 dtensor-opt                  0x0000562b94e82d17
13 dtensor-opt                  0x0000562b94e82ae5
14 dtensor-opt                  0x0000562b948df256
15 dtensor-opt                  0x0000562b948de53a
16 dtensor-opt                  0x0000562b9502d603
17 dtensor-opt                  0x0000562b9502d34e
18 dtensor-opt                  0x0000562b948da857
19 dtensor-opt                  0x0000562b948dad3d
20 dtensor-opt                  0x0000562b92d80fcc
21 libc.so.6                    0x00007f49de626083 __libc_start_main + 243
22 dtensor-opt                  0x0000562b92d80d39
/b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir:277:17: error: CHECK-LABEL: expected string not found in input
// CHECK-LABEL: module @test_spmd_softmax_rank_3
                ^
<stdin>:153:45: note: scanning from here
module @test_spmd_softmax_last_dim_unsharded {
                                            ^
<stdin>:156:303: note: possible intended match here
 %0 = "tf.Softmax"(%arg0) {_global_shape = [#tf_type.shape<32x32>], _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<16x32xf32>) -> tensor<16x32xf32>
                                                                                                                                                                                                                                                                                                              ^

Input file: <stdin>
Check file: /b/f/w/bazel-out/k8-opt/bin/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir.test.runfiles/org_tensorflow/tensorflow/dtensor/mlir/tests/spmd_expansion.mlir

-dump-input=help explains the following input dump.

Input was:
<<<<<<
             .
             .
             .
           148:  }
           149: }
           150:
           151:
           152: // -----
           153: module @test_spmd_softmax_last_dim_unsharded {
label:277'0                                                 X~~ error: no match found
           154:  func.func @main(%arg0: tensor<16x32xf32> {tf._global_shape = #tf_type.shape<32x32>, tf._layout = "sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"}) {
label:277'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           155:  "tf_device.cluster"() ({
label:277'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~
           156:  %0 = "tf.Softmax"(%arg0) {_global_shape = [#tf_type.shape<32x32>], _layout = ["sharding_specs:x,unsharded, mesh:TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"]} : (tensor<16x32xf32>) -> tensor<16x32xf32>
label:277'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
label:277'1                                                                                                                                                                                                                                                                                                                   ?                         possible intended match
           157:  tf_device.return {_global_shape = [], _layout = []}
label:277'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           158:  }) {_global_shape = [], _mesh = "TPU|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:TPU:0,/job:localhost/task:0/device:TPU:1,/job:localhost/task:0/device:TPU:2,/job:localhost/task:0/device:TPU:3"} : () -> ()
label:277'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           159:  return
label:277'0     ~~~~~~~~
           160:  }
label:277'0     ~~~
           161: }
label:277'0     ~~
           162:
label:277'0     ~
>>>>>>

--

********************
********************
Failed Tests (1):
  MLIR tests :: spmd_expansion.mlir
```
