func.func @main(%arg0: tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>) {
    %0 = "tf.Relu"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    %1 = "tf.Relu"(%0) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
    %2 = "tf.Add"(%0, %1) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %3 = "tf.Relu"(%2) : (tensor<?xi32>) -> tensor<?xi32>
    %4 = "tf.Relu"(%1) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>

// CHECK-LABEL: @main

      %0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
      return %0 : tensor<1x4x4x3xf32>
    }

  %cst = arith.constant dense<0.0> : tensor<2x3xbf16>
  %0 = "tfl.maximum"(%arg0, %cst) : (tensor<2x3xbf16>, tensor<2x3xbf16>) -> tensor<2x3xbf16>
  func.return %0 : tensor<2x3xbf16>

  // CHECK: %[[RESULT:.*]] = "tfl.relu"(%arg0)
  // CHECK: return %[[RESULT]]
}

// CHECK-LABEL: fuseScalarAddIntoConv2dBf16
func.func @fuseScalarAddIntoConv2dBf16(%arg0: tensor<256x32x32x3xbf16>, %arg1: tensor<16x3x3x3xbf16>) -> tensor<256x8x7x16xbf16> {

    func_op.eraseResult(0);
    func_op.insertResult(0, new_call_op.getResult(0).getType(),
                         /*resultAttrs=*/nullptr);

    // Modify the quantized fused function to do dequantize+relu(6).
    rewriter.setInsertionPoint(req_op);
    Value new_result = rewriter.create<mlir::stablehlo::UniformDequantizeOp>(
        req_op.getLoc(), func_op.getResultTypes()[0], req_op.getOperand());

      %outputs_12, %control_13 = tf_executor.island wraps "tf.Relu"(%outputs_10) {device = ""} : (tensor<16x112x112x?xf32>) -> tensor<16x112x112x?xf32>