} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
  %4 = "tf.BiasAdd"(%3, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
  %5 = "tf.Relu"(%4) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>


  %6 = "tf.Conv2D"(%arg0, %arg1) {
    data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [],

// CHECK-SAME: f = @composite_conv3d_fn_1}>
// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[RELU:.*]] = "tf.Relu"(%[[PARTITIONEDCALL_0]])
// CHECK: return %[[RELU]]

// CHECK-LABEL: private @composite_conv3d_fn_1

// WEIGHTONLY-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x3x2xf32>
// WEIGHTONLY: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]])

    parameters[
      {"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "!tf_type.qint8"},
      {"quantized_ops": ["${main_op}", "BiasAdd", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "!tf_type.qint8"},
      {"quantized_ops": ["${main_op}", "BiasAdd", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "!tf_type.qint8"},
    ]

// Canonicalize tf.Maximum of zero to tf.Relu
//===----------------------------------------------------------------------===//

def IsInteger32Pred: CPred<
  "getElementTypeOrSelf($0.getType()).isInteger(32)">;

// Whether the transformation is compatible with the device if given.
// Currently, Relu with int32 is not supported on GPU.
def IsDeviceCompatible: Constraint<

  // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Relu"], num_args = 1 : i64, operandSegmentSizes = array<i32: 1, 1, 1, 0>, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32>

      {"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "i8"},
      {"quantized_ops": ["${main_op}", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "i8"},
      {"quantized_ops": ["${main_op}"], "act_func": "internal_dequantize_no_activation_fn", "output_type": "f32"},

    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<1x5x5x2xf32>, tensor<3x1x1x2xf32>, tensor<3xf32>) -> tensor<1x5x5x3xf32>
  %conv2 = "tfl.conv_2d"(%0, %w, %b2) {
    dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU",
    padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32

  %none_attr = tfr.constant "NONE" -> !tfr.attr
  %relu_attr = tfr.constant "RELU" -> !tfr.attr
  %relu6_attr = tfr.constant "RELU6" -> !tfr.attr
  %reluN1_1_attr = tfr.constant "RELU_N1_TO_1" -> !tfr.attr
  %none:2 = "tfr.quant_act_range"(%none_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)
  %relu:2 = "tfr.quant_act_range"(%relu_attr, %scale, %zp) : (!tfr.attr, f32, i64) -> (!tfr.tensor, !tfr.tensor)

    // Currently, GPU only supports Conv2D+BiasAdd+Relu fusion.
    if (IsGpuDevice(conv)) {
      auto activation = GetActivation(bias_add);
      if (!activation || activation->getName().stripDialect() != "Relu" ||
          !bias_add.getOutput().hasOneUse()) {
        (void)rewriter.notifyMatchFailure(conv, [&](Diagnostic &diag) {
          diag << "GPU only supports Conv2D+BiasAdd+Relu fusion";
        });
        return false;

    %14 = "tfl.relu"(%10#1) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %15 = "tfl.logistic"(%10#0) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %16 = tfl.mul %15, %14 {fused_activation_function = "NONE"} : tensor<4x2xf32>
    %17 = tfl.add %13, %16 {fused_activation_function = "NONE"} : tensor<4x2xf32>
    %18 = "tfl.relu"(%17) : (tensor<4x2xf32>) -> tensor<4x2xf32>