%w = arith.constant dense<127.0> : tensor<3x3x3x3xf32>
  %b = arith.constant dense<0.0> : tensor<3xf32>
  %conv = "tfl.conv_2d"(%arg0, %w, %b) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x5x5x3xf32>, tensor<3x3x3x3xf32>, tensor<3xf32>) -> tensor<1x5x5x3xf32>
  func.return %conv : tensor<1x5x5x3xf32>

    %12 = "quantfork.qcast"(%11) {volatile} : (tensor<2x3x1x3xf32>) -> tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>
    %13 = "quantfork.dcast"(%12) : (tensor<2x3x1x3x!quant.uniform<i8:f32, 3.000000e-01:1>>) -> tensor<2x3x1x3xf32>
    return %13 : tensor<2x3x1x3xf32>
  }

  // CHECK: quantized_dot_general_fn_1

    return %5 : tensor<1x3x1x1xf32>
  }
  func.func private @composite_gather_fn_1(%arg0: tensor<1x3x1x1xf32>, %arg1: tensor<1xi32>, %arg2: tensor<i32>) -> tensor<1x3x1x1xf32> attributes {tf_quant.composite_function} {

    : (tensor<*x!tf_type.resource>) -> (tensor<3x3x1x32xf32>, tensor<3x3x1x32xf32>)
  "tf_device.launch"() ({
    "tf.TPUCompileSucceededAssert"(%compile#0) : (tensor<!tf_type.string>) -> ()
    tf_device.return
  }) {device = "/device:CPU:0"} : () -> ()
  %execute0 = "tf_device.launch"() ({
    %3 = "tf.TPUExecute"(%2#0, %2#1, %compile#1)
      : (tensor<3x3x1x32xf32>, tensor<3x3x1x32xf32>, tensor<2x!tf_type.string>) -> tensor<i32>

  %9 = "tfl.quantize"(%8) {qtype = tensor<1x3x1x1x!quant.uniform<i8:f32, 1.0:2>>, volatile} : (tensor<1x3x1x1xf32>) -> tensor<1x3x1x1x!quant.uniform<i8:f32, 1.0:2>>
  %10 = "tfl.dequantize"(%9) : (tensor<1x3x1x1x!quant.uniform<i8:f32, 1.0:2>>) -> tensor<1x3x1x1xf32>

  // CHECK-SAME: %[[ARG1:.*]]: tensor<*x!tf_type.resource<tensor<f32>>> {tf.device = "/device:TPU:1"},
  // CHECK-SAME: %[[ARG2:.*]]: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>> {tf.device = "/device:TPU:0"},
  // CHECK-SAME: %[[ARG3:.*]]: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>> {tf.device = "/device:TPU:1"})
  func.func @main(%arg0: !tf_res_f32 {tf.device = "/device:TPU:0"},
             %arg1: !tf_res_f32 {tf.device = "/device:TPU:1"},

// CHECK-LABEL: conv2d_backprop_input_with_add
func.func @conv2d_backprop_input_with_add(%arg0: tensor<4xi32>, %arg1: tensor<3x3x1x32xf32>, %arg2: tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> {
  %0 = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {strides = [1, 2, 2, 1], padding="SAME", dilations=[1, 1, 1, 1]}: (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32>

  // CHECK-SAME: (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>
  // CHECK: return %[[CALL]]

  func.func private @composite_dot_general_fn(%arg0: tensor<4x3x6x5xf32>, %arg1: tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : (tensor<4x3x6x5xf32>, tensor<4x3x5x2xf32>) -> tensor<4x3x6x2xf32>

        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>

    %9 = stablehlo.convert %3 : (tensor<3x3x4x4xi8>) -> tensor<3x3x4x4xf32>
    %10 = stablehlo.convolution(%8, %9) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
    %11 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>