: (tensor<?x3x2x1xi8>, tensor<2x1x1x1xi8>) -> tensor<?x2x2x1xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xi32>, tensor<?x2x2x1xi32>) -> tensor<?x2x2x1xi32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
  %3 = chlo.broadcast_add %2, %bias : (
      tensor<?x2x2x1xi32>, tensor<1xi32>) ->tensor<?x2x2x1xi32>
  return %3 : tensor<?x2x2x1xi32>
}

// -----

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

    }> {
      _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

  %2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
  return %2 : tensor<1x2x2x2xf32>
}
// CHECK-LABEL: @composite_conv2d_with_bias_and_relu6_fn_1
// CHECK-NOT: "tf.CalibrationStatisticsSaver"

// -----

// Check the IfOp is set to stateful.

    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>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

  %3 = "tfl.quantize"(%2) {qtype = tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>, volatile} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>
  %4 = "tfl.dequantize"(%3) : (tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>) -> tensor<1x2x2x3xf32>

    %1 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsit...

    rhs_quantization_max_val = 127 : i64
  } : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x2x1x2xf32>
  func.return %0 : tensor<1x2x1x2xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantized_convolution_hybrid_explicit
func.func @uniform_quantized_convolution_hybrid_explicit(%input: tensor<1x2x2x3xf32>) -> tensor<1x3x3x2xf32> {

module {
  func.func @conv_with_single_layer(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {