%2 = stablehlo.convolution(%1, %0)
      dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
      window = {
        stride = [1, 1], pad = [[0, 0], [1, 1]],
        lhs_dilate = [1, 1],
        rhs_dilate = [1, 1]
      }
      {
        batch_group_count = 1 : i64,
        feature_group_count = 1 : i64
      } : (tensor<1x3x2x3xf32>, tensor<2x3x3x2xf32>)

      %1 = "tf.XlaCallModule"(%0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_2, _original_entry_function = "composite_fn_2", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
      return %1 : tensor<1x4x4x3xf32>

}
// CHECK: @main(%[[ARG:.+]]: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32>

// Contains the `stablehlo.transpose` op of the arg (e.g. [b, f, 0, 1] to
// [b, 0, 1, f]). The weight constant is folded into [0, 1, i, o] format.
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<3.000000e+00> : tensor<3x3x8x8xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %arg0, dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>

    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    // CHECK-SKIP: tf.PartitionedCall
    // CHECK-NOSKIP: call @some_other_func
    // CHECK-NOSKIP-NOT: tf.PartitionedCall
    func.return %1: tensor<1x2x2x3xf32>
  }

  // CHECK-SKIP: func.func private @some_func
  func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {

// -----

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

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

func.func @testConv2DGPU(tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32> {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
  // CHECK: tac.cost = 0x4C300000

  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<[1, 3]> : tensor<2xi64>}> : (tensor<2x3xf32>) -> tensor<1x2x2x3xf32>
  %1 = mhlo.multiply %0, %cst1 : tensor<1x2x2x3xf32>
  // CHECK:      return %[[RES]] : tensor<1x2x2x3xf32>
  func.return %1 : tensor<1x2x2x3xf32>
}

// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_1D_int

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

// CHECK: return %[[RES0]] : tensor<256x32x32x16xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "CONV_2D"
    }
  ],
  "subgraphs": [
    {

// CHECK: return %[[RES0]] : tensor<256x32x32x16xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "CONV_2D"
    }
  ],
  "subgraphs": [