// CHECK: %cst_0 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %2 = "tfl.transpose"(%1, %cst_0) : (tensor<1x1x4x1xf32>, tensor<4xi32>) -> tensor<1x1x1x4xf32>
// CHECK: return %2 : tensor<1x1x1x4xf32>


func.func @avg_pool2d_5(%arg0: tensor<1x1x3x3xf32>) -> (tensor<1x1x2x2xf32>) {

    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    func.return %0: tensor<1x2x2x3xf32>
  }

  func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> {
    return %arg0 : tensor<1x2x2x3xf32>
  }
}

// CHECK: module
// CHECK-NOT: tf.PartitionedCall
// CHECK-NOT: some_func

func.func @dont_move_transpose_different_ranks(%arg0:tensor<1x1x2x3xf32>, %arg1:tensor<2x3xf32>) -> tensor<1x2x1x3xf32> {
  %cst = "tf.Const"() {value = dense<[0, 2, 1, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
  %0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<1x1x2x3xf32>, tensor<2x3xf32>) -> tensor<1x1x2x3xf32>
  %1 = "tf.Transpose"(%0, %cst) {device = ""} : (tensor<1x1x2x3xf32>, tensor<4xi32>) -> tensor<1x2x1x3xf32>

  func.return %1 : tensor<1x2x1x3xf32>

  %mm = "tfl.batch_matmul"(%w, %0) {adj_x = false, adj_y = true} : (tensor<1x1x12x512xf32>, tensor<1x1x3x512xf32>) -> tensor<1x1x12x3xf32>
  %mm_s = "quantfork.stats"(%mm) {layerStats = dense<[0.000000e+00, 1.000000e+01]> : tensor<2xf32>} : (tensor<1x1x12x3xf32>) -> tensor<1x1x12x3xf32>
  func.return %mm_s : tensor<1x1x12x3xf32>

// CHECK: %[[w:.*]] = arith.constant dense<1.270000e+02> : tensor<1x1x12x512xf32>

// CHECK:  [[VAL_1:%.*]] = "tfl.reshape"(%2, %[[CST]]) {tac.device = "GPU",  tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:  [[VAL_2:%.*]] = "tfl.concatenation"([[VAL_0]], [[VAL_1]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>

}

// CHECK-LABEL: conv_with_relu1_pattern1
func.func @conv_with_relu1_pattern1(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x4x2xf32>) {
  %cst = "tf.Const"() {value = dense<[[[[-8.69931221, 6.44628429], [-9.18393421, 1.53671741], [8.68561744, -3.581774]]]]> : tensor<1x1x3x2xf32>} : () -> tensor<1x1x3x2xf32>
  %cst_0 = "tf.Const"() {value = dense<-1.000000e+00> : tensor<f32>} : () -> tensor<f32>

  func.return %0 : tensor<256x32x32x16xf32>
}

func.func @testConv2DDynamicShape(tensor<?x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<?x32x32x16xf32> {
^bb0(%arg0: tensor<?x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
  // CHECK: _arithmetic_count = -1 : i64

func.func @testConv(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>

    %18 = stablehlo.multiply %16, %17 : tensor<1x3x3x4xf32>
    %19 = call @uniform_quantize_1(%18, %5, %6) : (tensor<1x3x3x4xf32>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xi8>
    %20 = call @uniform_dequantize_0(%19, %5, %6) : (tensor<1x3x3x4xi8>, tensor<1x1x1x1xf32>, tensor<1x1x1x1xi8>) -> tensor<1x3x3x4xf32>
    return %20 : tensor<1x3x3x4xf32>
  }

    func.return %1: tensor<*xf32>
  }

  func.func private @outer_fn(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x1024xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.PartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @inner_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>