{"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"},
      {"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_dequantize_and_relu_fn", "output_type": "f32"},

  %1 = "tfl.relu6"(%0) : (tensor<256x30x30x16xf32>) -> tensor<256x30x30x16xf32>
  func.return %1 : tensor<256x30x30x16xf32>

  %biasadd = "tf.BiasAdd"(%conv, %dq_bias) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
  %res = "tf.Relu6"(%biasadd) : (tensor<*xf32>) -> tensor<*xf32>
  %q_res = "quantfork.qcast"(%res) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.023529411764705882:-128>>

    %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %2 = "tf.Relu6"(%1) : (tensor<*xf32>) -> tensor<*xf32>
    func.return %2 : tensor<*xf32>
  }

    %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %2 = "tf.Relu6"(%1) : (tensor<*xf32>) -> tensor<*xf32>
    func.return %2 : tensor<*xf32>
  }

// CHECK-LABEL: func @conv_with_single_layer

// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0
// CHECK-SAME: f = @quantize_i8

  %biasadd = "tf.BiasAdd"(%conv, %dq_bias) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
  %res = "tf.Relu6"(%biasadd) : (tensor<*xf32>) -> tensor<*xf32>
  %q_res = "quantfork.qcast"(%res) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.023529411764705882:-128>>

// CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[dq_input]], %[[dq_weight]])
// CHECK: %[[bias:.*]] = "tf.BiasAdd"(%[[conv]], %[[cst_0]]) <{data_format = "NHWC"}>
// CHECK: %[[relu6:.*]] = "tf.Relu6"(%[[bias]])

// -----

func.func @remove_check_numerics_op(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  %0 = "tf.CheckNumerics"(%arg0) {device = "", message = "transformer"} : (tensor<*xf32>) -> tensor<*xf32>

    [(IsNotInLiftedFunc $res),
     (IsEinsumSupportedByXlaDotV2 $equation)],
    [], (addBenefit 10)>;

}
defm : LiftCompositeOpsWithActivation<TF_ReluOp, "relu">;
defm : LiftCompositeOpsWithActivation<TF_Relu6Op, "relu6">;

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">
    (ArgumentList $params, $indices, $axis),
    (ResultList $res),

  [(CopyAttrs $maximum_op, $dest)]>;

//===----------------------------------------------------------------------===//
// Canonicalize tf.Relu of Minimul six to tf.Relu6
//===----------------------------------------------------------------------===//

def ReluOfMinimum6ToRelu6 : Pat<
  (TF_ReluOp:$src (TF_MinimumOp $x, $y)), (TF_Relu6Op:$dest $x),

  RELU = 19,
  // NOTE(aselle): RELU_N1_TO_1 used to be called RELU1, but it was renamed
  // since different model developers use RELU1 in different ways. Never
  // create another op called RELU1.
  RELU_N1_TO_1 = 20,
  RELU6 = 21,
  RESHAPE = 22,
  RESIZE_BILINEAR = 23,
  RNN = 24,
  SOFTMAX = 25,
  SPACE_TO_DEPTH = 26,
  SVDF = 27,
  TANH = 28,
  CONCAT_EMBEDDINGS = 29,
  SKIP_GRAM = 30,
  CALL = 31,