//tensorflow/compiler/mlir/lite/tests:optimize.mlir.test fails
@cfRod @nSircombe @TensorFlow-MKL @milpuz01  Roll back https://github.com/tensorflow/tensorflow/commit/973cd1af48a7006a06bfbdbebbb659ad566d6b66 This is no longer a problem. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61440">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61440">No</a>

### 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.17

### Bazel version

6.1.0

### GCC/compiler version

10.2.1

### CUDA/cuDNN version

n/a

### GPU model and memory

n/a

### Current behavior?

Unit test fails since commit https://github.com/tensorflow/tensorflow/commit/9d0fea2d5935285122b56867c4499433121f531f

### Standalone code to reproduce the issue

```shell
bazel test --config=mkl_aarch64_threadpool --copt=-flax-vector-conversions --test_env=TF_ENABLE_ONEDNN_OPTS=1 --test_env=TF2_BEHAVIOR=1 --define=tf_api_version=2 --jobs=75 --build_tests_only -- //tensorflow/compiler/mlir/lite/tests:optimize.mlir.test
```


### Relevant log output

```shell
==================== Test output for //tensorflow/compiler/mlir/lite/tests:optimize.mlir.test:
-- Testing: 1 tests, 1 workers --
FAIL: MLIR tests :: optimize.mlir (1 of 1)
******************** TEST 'MLIR tests :: optimize.mlir' FAILED ********************
Script:
--
: 'RUN: at line 2';   /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/tf-opt /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir -tfl-optimize | /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/llvm-project/llvm/FileCheck /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir
: 'RUN: at line 4';   /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/tf-opt /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir -tfl-optimize='enable-canonicalization=true' | /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/llvm-project/llvm/FileCheck --check-prefix=FOLD /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir
: 'RUN: at line 7';   /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/tf-opt /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir -tfl-legalize-tf -tfl-optimize | /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/llvm-project/llvm/FileCheck --check-prefix=Fusing /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir
: 'RUN: at line 9';   /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/tf-opt /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir -tfl-legalize-tf -tfl-optimize='disable-fuse-mul-and-fc=true' | /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/llvm-project/llvm/FileCheck --check-prefix=NoFusing /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir
--
Exit Code: 1

Command Output (stderr):
--
/home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir:3622:12: error: CHECK: expected string not found in input
 // CHECK: "tfl.batch_matmul"(%arg0, %arg1)
           ^
<stdin>:1542:53: note: scanning from here
 func.func @FuseReshapeAndTransposeAroundBatchMatmul(%arg0: tensor<1x128x1024xf32>, %arg1: tensor<1024x16xf32>) -> tensor<1x128x16xf32> {
                                                    ^
<stdin>:1549:7: note: possible intended match here
 %2 = "tfl.batch_matmul"(%arg1, %1) {adj_x = true, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<1024x16xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
      ^

Input file: <stdin>
Check file: /home/andrew/src/tf_test/tensorflow-git/bazel-ci_build-cache/.cache/bazel/_bazel_andrew/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/tensorflow/compiler/mlir/lite/tests/optimize.mlir.test.runfiles/org_tensorflow/tensorflow/compiler/mlir/lite/tests/optimize.mlir

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

Input was:
<<<<<<
              .
              .
              .
           1537:  %3 = "tfl.transpose"(%2, %0) : (tensor<16x1x8x1280xf32>, tensor<4xi32>) -> tensor<8x1x16x1280xf32> 
           1538:  %4 = "tfl.gather_nd"(%3, %1) : (tensor<8x1x16x1280xf32>, tensor<16x1xi32>) -> tensor<16x1x16x1280xf32> 
           1539:  %5 = "tfl.reshape"(%4, %cst) : (tensor<16x1x16x1280xf32>, tensor<4xi32>) -> tensor<1x16x16x1280xf32> 
           1540:  return %5 : tensor<1x16x16x1280xf32> 
           1541:  } 
           1542:  func.func @FuseReshapeAndTransposeAroundBatchMatmul(%arg0: tensor<1x128x1024xf32>, %arg1: tensor<1024x16xf32>) -> tensor<1x128x16xf32> { 
check:3622'0                                                         X~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ error: no match found
           1543:  %cst = arith.constant dense<[1, 2, 0]> : tensor<3xi32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1544:  %cst_0 = arith.constant dense<[16, 1, 128]> : tensor<3xi32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1545:  %cst_1 = arith.constant dense<[1024, 128]> : tensor<2xi32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1546:  %cst_2 = arith.constant dense<[2, 0, 1]> : tensor<3xi32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1547:  %0 = "tfl.transpose"(%arg0, %cst_2) : (tensor<1x128x1024xf32>, tensor<3xi32>) -> tensor<1024x1x128xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1548:  %1 = "tfl.reshape"(%0, %cst_1) : (tensor<1024x1x128xf32>, tensor<2xi32>) -> tensor<1024x128xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1549:  %2 = "tfl.batch_matmul"(%arg1, %1) {adj_x = true, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<1024x16xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
check:3622'1           ?                                                                                                                                                                    possible intended match
           1550:  %3 = "tfl.reshape"(%2, %cst_0) : (tensor<16x128xf32>, tensor<3xi32>) -> tensor<16x1x128xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1551:  %4 = "tfl.transpose"(%3, %cst) : (tensor<16x1x128xf32>, tensor<3xi32>) -> tensor<1x128x16xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1552:  return %4 : tensor<1x128x16xf32> 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
           1553:  } 
check:3622'0     ~~~
           1554:  func.func @FuseTransposeFCRhsToBatchMatmul(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xf32>, %arg2: none) -> tensor<16x128xf32> { 
check:3622'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
              .
              .
              .
>>>>>>

--

********************
********************
Failed Tests (1):
  MLIR tests :: optimize.mlir


Testing Time: 0.53s
  Failed: 1
================================================================================
```
