%0 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x2x4xf32>
  %1 = stablehlo.transpose %arg0, dims = [0, 2, 1, 3] : (tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32>
  %2 = stablehlo.add %1, %0 : tensor<1x3x2x4xf32>
  return %2 : tensor<1x3x2x4xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose
// CHECK: %[[ADD_0:.+]] = stablehlo.add %[[TRANSPOSE_0]], {{.*}}
// CHECK: return %[[ADD_0]]

// -----

// CHECK{LITERAL}: %cst_1 = arith.constant dense<[[[[1.000000e+00], [2.000000e+00]], [[2.000000e+00], [4.000000e+00]]]]> : tensor<1x2x2x1xf32>
// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x2x2x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %4 = "tfl.transpose"(%3, %cst_2) : (tensor<1x2x2x1xf32>, tensor<4xi32>) -> tensor<1x1x2x2xf32>
// CHECK: return %4 : tensor<1x1x2x2xf32>

// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>
// CHECK-DAG: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[BROADCAST]] : tensor<?x3x2x2xf32>
// CHECK: return %[[ADD]] : tensor<?x3x2x2xf32>

// -----

func.func @func_conv_batchnorm_relu6(%arg_0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {

// CHECK: return %[[identity]] : tensor<1x3x2x2xf32>

func.func @cast_bf16_avg_pool_to_fp32(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {
  %cst = "tf.Const"() {device = "", value = dense<1.000000e+00> : tensor<2x3x3x2xbf16>} : () -> tensor<2x3x3x2xbf16>
  %0 = "tf.Cast"(%arg0) {Truncate = false, device = ""} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xbf16>

  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  // CHECK: %[[MUL:.*]] = mhlo.multiply %[[BROADCAST]], %[[ARG]] : tensor<1x1x2x4xi32>
  %1 = mhlo.multiply %0, %arg0 : tensor<1x1x2x4xi32>
  // CHECK:      return %[[MUL]] : tensor<1x1x2x4xi32>
  func.return %1 : tensor<1x1x2x4xi32>

    %3 = "tf.Identity"(%2) {device = ""} : (tensor<1x3x1x1xf32>) -> tensor<1x3x1x1xf32>
    return %3 : tensor<1x3x1x1xf32>
  }

// CHECK-LABEL: func @multiple_quantizable_ops_in_graph

    %2 = "tf.Relu"(%1) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
    %3 = "tf.FakeQuantWithMinMaxArgs"(%2) {device = "", max = 4.000000e-01 : f32, min = -3.000000e-01 : f32, narrow_range = false, num_bits = 8 : i64} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>

    has_token_input_output = false, module = "", platforms = [],
    version = 5 : i64
  } : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
  return %0 : tensor<1x2x2x2xf32>
}

// CHECK-LABEL: func.func @qdq_for_conv_weight_empty
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>

func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>

    %0 = "tf.Cast"(%arg0) : (tensor<1x3x2x2xi8>) -> tensor<1x3x2x2xi32>
    %1 = "tf.Sub"(%0, %arg2) : (tensor<1x3x2x2xi32>, tensor<i32>) -> tensor<1x3x2x2xi32>
    %2 = "tf.Cast"(%1) : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32>
    %3 = "tf.Mul"(%2, %arg1) : (tensor<1x3x2x2xf32>, tensor<f32>) -> tensor<1x3x2x2xf32>
    return %3 : tensor<1x3x2x2xf32>
  }