return %0 : tensor<?x768x3xf32>
// CHECK:  %cst = arith.constant dense<[-9223372036854775808, 768, 2, 1]> : tensor<4xi64>
// CHECK:  %0 = "tf.Reshape"(%arg0, %cst) : (tensor<?x768x2xf32>, tensor<4xi64>) -> tensor<?x768x2x1xf32>
// CHECK:  %cst_0 = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : () -> tensor<4xi64>
// CHECK:  %1 = "tf.Transpose"(%0, %cst_0) : (tensor<?x768x2x1xf32>, tensor<4xi64>) -> tensor<?x2x1x768xf32>

  func.return %0 : tensor<40x40xf32>
// CHECK-LABEL: matmul
// CHECK: %[[CST:.*]] = arith.constant dense<[1, 0]> : tensor<2xi32>
// CHECK: %[[ARG:.*]] = "tfl.transpose"(%arg1, %[[CST]]) : (tensor<37x40xf32>, tensor<2xi32>) -> tensor<40x37xf32>
// CHECK: %[[CST_0:.*]] = "tfl.no_value"() <{value}> : () -> none

  %1 = "tfl.dequantize"(%0) : (tensor<1x32x1x3x!quant.uniform<i8:f32, 1.0>>) -> tensor<1x32x1x3xf32>
  %cst = arith.constant dense<1> : tensor<4xi32>
  %cst_0 = arith.constant dense<[0, 0, 0, 3]> : tensor<4xi32>
  %cst_1 = arith.constant dense<[0, -2, 0, 0]> : tensor<4xi32>

    moved after _TPUCompileMlir and the first input will use the
    _TPUCompileMlir program output:

    ```mlir
    "tf_device.launch"() ({
     %cst_0 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
     "tf.OpA"(%cst_0) { mini_batch_splits = ""} : (tensor<1x!tf_type.string>) -> ()
     tf_device.return
   }) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ()

  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} : () -> tensor<10xf32>
  %cst_0 = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi32>} : () -> tensor<2xi32>
  %1 = "tf.MatMul"(%arg0, %arg1) {
    transpose_a = false, transpose_b = true
  } : (tensor<1x10xf32>, tensor<10x10xf32>) -> tensor<*xf32>
  %2 = "tf.Reshape"(%1, %cst_0) : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>

// QDQ-NEXT:  return %[[split]]#0 : tensor<2xf32>
}

// CHECK-LABEL: FoldTranspose
func.func @FoldTranspose(%arg0: tensor<1x10x20x3xf32>) -> tensor<1x20x40x16xf32> {
  %cst = arith.constant dense<[1, 20, 40, 16]> : tensor<4xi32>
  %cst_0 = arith.constant dense<[2, 0, 1, 3]> : tensor<4xi32>

func.func @FuseTransposeFCLhsToBatchMatmul(%arg0: tensor<1024x4xf32>, %arg1: tensor<8x1024xf32>, %arg2: tensor<4x256xf32>) -> tensor<8x256xf32> {
  %cst_0 = arith.constant dense<[1, 0]> : tensor<2xi32>
  %cst_1 = "tfl.no_value"() {value} : () -> none
  %0 = "tfl.transpose"(%arg0, %cst_0) : (tensor<1024x4xf32>, tensor<2xi32>) -> tensor<4x1024xf32>

constexpr absl::string_view kModuleConstantAdd = R"mlir(
  module {
    func.func @constant_add() -> (tensor<3x2xf32>) {
      %cst1 = stablehlo.constant dense<2.4> : tensor<3x2xf32>
      %cst2 = stablehlo.constant dense<5.7> : tensor<3x2xf32>
      %add = stablehlo.add %cst1, %cst2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
      func.return %add : tensor<3x2xf32>
    }
  }
)mlir";

func.func @decompose_batch_norm(%arg0: tensor<*xf32>) -> (tensor<*xf32>) {
  %cst = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32>
  %cst_0 = "tf.Const"() {value = dense<0.500000e+00> : tensor<2xf32>} : () -> tensor<2xf32>

module {
  func.func @restore_single() {
    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>